Merge pull request #709 from plasmashen/master

Fix the problem of text output from T5 model
QData · Aug 28, 2023 · e3d05c4 · e3d05c4
2 parents 2abb4dd + 25f9b78
commit e3d05c4
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/textattack/goal_functions/goal_function.py b/textattack/goal_functions/goal_function.py
@@ -176,13 +176,15 @@ def _call_model_uncached(self, attacked_text_list):
             if isinstance(batch_preds, list):
                 outputs.extend(batch_preds)
             elif isinstance(batch_preds, np.ndarray):
-                outputs.append(torch.tensor(batch_preds))
+                outputs.append(batch_preds)
             else:
                 outputs.append(batch_preds)
             i += self.batch_size
 
         if isinstance(outputs[0], torch.Tensor):
             outputs = torch.cat(outputs, dim=0)
+        elif isinstance(outputs[0], np.ndarray):
+            outputs = np.concatenate(outputs).ravel()
 
         assert len(inputs) == len(
             outputs