Added proper cuda support for mask

experimenti · Jul 12, 2018 · c959d29 · c959d29
1 parent 78e1607
commit c959d29
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/model.py b/model.py
@@ -58,7 +58,9 @@ def forward(self, x, lengths):
  """
  for module in self.seq_module:
  x = module(x)
- mask = torch.ByteTensor(x.size()).fill_(0).cuda() # TODO don't hard-code cuda
+ mask = torch.ByteTensor(x.size()).fill_(0)
+ if x.is_cuda:
+ mask = mask.cuda()
  for i, length in enumerate(lengths):
  length = length.item()
  if (mask[i].size(2) - length) > 0: