removing tensor resizing in future_mask
Summary:
tensor resizing doesn't work well with tpus, this change is equivalent
to the base and works better w/ tpus.
Pull Request resolved:
@@ -462,10 +462,8 @@ class TransformerDecoder(FairseqIncrementalDecoder):
def buffered_future_mask(self, tensor):
dim = tensor.size(0)
- if not hasattr(self, '_future_mask') or self._future_mask is None or self._future_mask.device != tensor.device:
+ if not hasattr(self, '_future_mask') or self._future_mask is None or self._future_mask.device != tensor.device or self._future_mask.size(0) < dim:
self._future_mask = torch.triu(utils.fill_with_neg_inf(tensor.new(dim, dim)), 1)
- if self._future_mask.size(0) < dim:
- self._future_mask = torch.triu(utils.fill_with_neg_inf(self._future_mask.resize_(dim, dim)), 1)
return self._future_mask[:dim, :dim]
def upgrade_state_dict_named(self, state_dict, name):
