utils.py: compatibility with new pytorch

NVIDIA · Dec 12, 2019 · a513db5 · a513db5
1 parent 37a033d
commit a513db5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/utils.py b/utils.py
@@ -6,7 +6,7 @@
 def get_mask_from_lengths(lengths):
     max_len = torch.max(lengths).item()
     ids = torch.arange(0, max_len, out=torch.cuda.LongTensor(max_len))
-    mask = (ids < lengths.unsqueeze(1)).byte()
+    mask = (ids < lengths.unsqueeze(1)).bool()
     return mask