fix hidden states and scores not being returned

neulab · Jul 9, 2023 · ccded32 · ccded32
1 parent fcc523d
commit ccded32
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/llm_client/__init__.py b/llm_client/__init__.py
@@ -55,7 +55,7 @@ def prompt(
         Returns:
         """
         if isinstance(text, str):
-            return self.prompt([text], max_new_tokens, **kwargs)
+            return self.prompt([text], max_new_tokens, output_scores, output_hidden_states, **kwargs)
 
         # TODO: Check for max length limit to avoid OOMs