X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?p=pytorch.git;a=blobdiff_plain;f=attentiontoy1d.py;h=2cecad8b18c9fb14e779a3bb517da0eaeceefecf;hp=e82894e650a05c36b8e6779dcf840fa3182319d7;hb=de3a0375a79cbbf4299aacc41db1426a39ca9664;hpb=d2c145b4306d5c36094618ff7e7323c5d083e1df

diff --git a/attentiontoy1d.py b/attentiontoy1d.py
index e82894e..2cecad8 100755
--- 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(