Fixed bug in testing script, swap to BoolTensor for PyTorch 1.2 support

SeanNaren · Aug 26, 2019 · ea7b424 · ea7b424
1 parent e73ccf6
commit ea7b424
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 6 deletions.
diff --git a/model.py b/model.py
@@ -57,7 +57,7 @@ def forward(self, x, lengths):
         """
         for module in self.seq_module:
             x = module(x)
-            mask = torch.ByteTensor(x.size()).fill_(0)
+            mask = torch.BoolTensor(x.size()).fill_(0)
             if x.is_cuda:
                 mask = mask.cuda()
             for i, length in enumerate(lengths):

diff --git a/test.py b/test.py
@@ -39,14 +39,10 @@ def evaluate(test_loader, device, model, decoder, target_decoder, save_output=Fa
 
         out, output_sizes = model(inputs, input_sizes)
 
-        if save_output:
-            # add output to data array, and continue
-            output_data.append((out.cpu().numpy(), output_sizes.numpy()))
-
         decoded_output, _ = decoder.decode(out, output_sizes)
         target_strings = target_decoder.convert_to_strings(split_targets)
 
-        if args.save_output is not None:
+        if save_output is not None:
             # add output to data array, and continue
             output_data.append((out.cpu().numpy(), output_sizes.numpy(), target_strings))
         for x in range(len(target_strings)):