From 5743fae23fe419197ca80bfa2492d3a3f5b7cbbe Mon Sep 17 00:00:00 2001
From: vince62s <vince62s@yahoo.com>
Date: Fri, 20 Oct 2023 15:21:11 +0200
Subject: [PATCH] less important but align to beam_search

---
 onmt/translate/greedy_search.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/onmt/translate/greedy_search.py b/onmt/translate/greedy_search.py
index 3b402cc9e8..f2272d3c3d 100644
--- 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)