need offload_folder argument for HFLM initializer

Stability-AI · Oct 18, 2023 · 5e80a0b · 5e80a0b
1 parent 4533dfb
commit 5e80a0b
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/lm_eval/models/gpt2.py b/lm_eval/models/gpt2.py
@@ -13,6 +13,7 @@ def __init__(
         low_cpu_mem_usage=None,
         torch_dtype=None,
         device_map=None,
+        offload_folder=None,
         subfolder=None,
         tokenizer=None,
         batch_size=1,
@@ -49,6 +50,7 @@ def __init__(
             low_cpu_mem_usage=low_cpu_mem_usage,
             torch_dtype=torch_dtype,
             device_map=device_map,
+            offload_folder=offload_folder,
             revision=revision,
             trust_remote_code=trust_remote_code,
         ).eval()