fix

ku-nlp · Aug 8, 2023 · 34ddade · 34ddade
1 parent 3073e76
commit 34ddade
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/kwja/modules/seq2seq.py b/src/kwja/modules/seq2seq.py
@@ -103,7 +103,11 @@ def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> A
             else None,
             **self.hparams.decoding,
         )
+        if isinstance(generations, torch.Tensor):
+            seq2seq_predictions = generations
+        else:
+            seq2seq_predictions = generations["sequences"]
         return {
             "example_ids": batch["example_ids"] if "example_ids" in batch else [],
-            "seq2seq_predictions": generations,
+            "seq2seq_predictions": seq2seq_predictions,
         }