Load from revision (#292)

allenai · Aug 26, 2024 · 8445178 · 8445178
1 parent c82b0c6
commit 8445178
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/eval/utils.py b/eval/utils.py
@@ -251,6 +251,7 @@ def load_hf_lm(
         if device_map:
             model = AutoModelForCausalLM.from_pretrained(
                 model_name_or_path,
+                revision=revision,
                 device_map=device_map,
                 torch_dtype=torch_dtype,
                 token=token,
@@ -259,6 +260,7 @@ def load_hf_lm(
         else:
             model = AutoModelForCausalLM.from_pretrained(
                 model_name_or_path,
+                revision=revision,
                 torch_dtype=torch_dtype,
                 token=token,
                 trust_remote_code=trust_remote_code,