Update.

[pytorch.git] / attentiontoy1d.py

diff --git a/attentiontoy1d.py b/attentiontoy1d.py
index e82894e..2cecad8 100755 (executable)
--- a/attentiontoy1d.py
+++ b/attentiontoy1d.py
@@ -321,6 +321,8 @@ if args.with_attention:
 test_input = test_input.detach().to('cpu')
 test_outputs = test_outputs.detach().to('cpu')
 test_targets = test_targets.detach().to('cpu')
+test_bx = test_bx.detach().to('cpu')
+test_tr = test_tr.detach().to('cpu')
 
 for k in range(15):
     save_sequence_images(