From de3a0375a79cbbf4299aacc41db1426a39ca9664 Mon Sep 17 00:00:00 2001 From: Francois Fleuret Date: Fri, 27 Nov 2020 23:13:08 +0100 Subject: [PATCH] Update. --- attentiontoy1d.py | 2 ++ 1 file changed, 2 insertions(+) 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( -- 2.20.1