Merge pull request #55 from huggingface/fix-trl-llava

Fix types to allow nullables in `llava_hf.py`
EvolvingLMMs-Lab · Apr 12, 2024 · 8519075 · 8519075
2 parents a876169 + 639f90c
commit 8519075
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/lmms_eval/models/llava_hf.py b/lmms_eval/models/llava_hf.py
@@ -46,9 +46,9 @@ def __init__(
         revision: str = "main",
         device: str = "cuda",
         dtype: Optional[Union[str, torch.dtype]] = "auto",
-        batch_size: Union[int, str] = 1,
+        batch_size: int = 1,
         trust_remote_code: Optional[bool] = False,
-        attn_implementation: Optional[str] = "flash_attention_2",
+        attn_implementation: Optional[str] = None,
         device_map: str = "",
         chat_template: Optional[str] = None,
         use_cache: bool = True,