fix device (#4887)

facebookresearch · Nov 21, 2022 · bbdd64b · bbdd64b
1 parent 4ffb29c
commit bbdd64b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/parlai/core/torch_generator_agent.py b/parlai/core/torch_generator_agent.py
@@ -1596,7 +1596,7 @@ def advance(self, logprobs, step):
             hyp_device = self.partial_hyps.get_device()
         self.partial_hyps = torch.cat(
             (
-                self.partial_hyps[path_selection.hypothesis_ids.long()],
+                self.partial_hyps[path_selection.hypothesis_ids.long().to(hyp_device)],
                 path_selection.token_ids.view(path_selection.token_ids.shape[0], -1).to(
                     hyp_device
                 ),