less important but align to beam_search

vince62s · Oct 20, 2023 · 5743fae · 5743fae
1 parent 7a2d8a4
commit 5743fae
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/onmt/translate/greedy_search.py b/onmt/translate/greedy_search.py
@@ -271,7 +271,7 @@ def update_finished(self):
                 best_hyp = sorted(self.hypotheses[b], key=lambda x: x[0], reverse=True)
                 for score, pred, attn in best_hyp:
                     self.scores[b].append(score)
-                    self.predictions[b].append(pred)
+                    self.predictions[b].append(pred.cpu())
                     self.attention[b].append(attn)
             return
         is_alive = ~self.is_finished.view(-1)