batch_size may be str if 'auto' is specified (#2084)

30273b47 · meg · GitHub · 058cfd0e · 30273b47
Unverified Commit 30273b47 authored Jul 10, 2024 by meg Committed by GitHub Jul 10, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

lm_eval/evaluator.py lm_eval/evaluator.py +1 -1

No files found.
--- a/lm_eval/evaluator.py
+++ b/lm_eval/evaluator.py
@@ -50,7 +50,7 @@ def simple_evaluate(
    model_args: Optional[Union[str, dict]] = None,
    tasks: Optional[List[Union[str, dict, object]]] = None,
    num_fewshot: Optional[int] = None,
-    batch_size: Optional[int] = None,
+    batch_size: Optional[Union[int, str]] = None,
    max_batch_size: Optional[int] = None,
    device: Optional[str] = None,
    use_cache: Optional[str] = None,