Browse Source

utils.py: compatibility with new pytorch

master
Rafael Valle 5 years ago
parent
commit
a513db50d0
1 changed files with 1 additions and 1 deletions
  1. +1
    -1
      utils.py

+ 1
- 1
utils.py View File

@ -6,7 +6,7 @@ import torch
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

Loading…
Cancel
Save