Update beam search code to support torch.bool change

Summary: Pull Request resolved: fairinternal/fairseq-py#797 Differential Revision: D16617067 Pulled By: myleott fbshipit-source-id: 52e3aeb98d6e3b55ff9154b784028bf13eabfe38
LLL-Orleans · Aug 2, 2019 · 5f34252 · 5f34252
1 parent ccb5dea
commit 5f34252
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/fairseq/search.py b/fairseq/search.py
@@ -212,7 +212,7 @@ def _sample_topp(self, lprobs):
 
         # trim the words that are not in top-P by setting their probabilities
         # to 0, so that they would not be sampled later.
-        trim_mask = truncated_mask.bitwise_not()
+        trim_mask = (~truncated_mask)
         trimed_probs = truncated_probs.masked_fill_(trim_mask, 0)
         return trimed_probs, truncated_indices
 

diff --git a/fairseq/sequence_generator.py b/fairseq/sequence_generator.py
@@ -163,7 +163,7 @@ def generate(
         # For example, suppose we're sampling and have already finalized 2/5
         # samples. Then the blacklist would mark 2 positions as being ignored,
         # so that we only finalize the remaining 3 samples.
-        blacklist = src_tokens.new(bsz, beam_size).byte().fill_(0)
+        blacklist = src_tokens.new_zeros(bsz, beam_size).eq(-1)  # forward and backward-compatible False mask
 
         # list of completed sentences
         finalized = [[] for i in range(bsz)]