fix type hint in PromptPipeline.__init__ (#496)

CarperAI · Jun 8, 2023 · 404217b · 404217b
1 parent 0dce99d
commit 404217b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/trlx/pipeline/offline_pipeline.py b/trlx/pipeline/offline_pipeline.py
@@ -123,7 +123,7 @@ class PromptPipeline(BasePipeline):
 
     def __init__(
         self,
-        prompts: Union[Dict[str, Any], List[str]],
+        prompts: Union[List[Dict[str, Any]], List[str]],
         max_prompt_length: int,
         tokenizer: PreTrainedTokenizer,
         add_special_tokens: bool = False,