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