Commit 2a681d99 authored by Alexei Baevski's avatar Alexei Baevski Committed by Myle Ott
Browse files

make sure tensor used to index is cuda if on gpu

parent 88df72c0

fairseq/sequence_generator.py

+1 −1
Changes for fairseq/sequence_generator.py: 1 added line, 1 removed line.
Original line number Diff line number Diff line
@@ -380,7 +380,7 @@ class SequenceGenerator(object): 
                new_bsz = bsz - len(finalized_sents)



                batch_mask = torch.ones(bsz).type_as(cand_indices)

                batch_mask[torch.LongTensor(finalized_sents)] = 0

                batch_mask[cand_indices.new(finalized_sents)] = 0

                batch_idxs = batch_mask.nonzero().squeeze(-1)



                eos_mask = eos_mask[batch_idxs]