Commit 9142999e authored by Benjamin Fattori's avatar Benjamin Fattori
Browse files

update type arguments for hf-causal to accept str for batch_size

parent e87c0838
...@@ -69,7 +69,7 @@ class HuggingFaceAutoLM(BaseLM): ...@@ -69,7 +69,7 @@ class HuggingFaceAutoLM(BaseLM):
tokenizer: Optional[str] = None, tokenizer: Optional[str] = None,
subfolder: Optional[str] = None, subfolder: Optional[str] = None,
revision: Optional[str] = "main", revision: Optional[str] = "main",
batch_size: Optional[int] = 1, batch_size: Optional[Union[int,str]] = 1,
max_gen_toks: Optional[int] = 256, max_gen_toks: Optional[int] = 256,
max_length: Optional[int] = None, max_length: Optional[int] = None,
add_special_tokens: Optional[bool] = None, add_special_tokens: Optional[bool] = None,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment