move hf model to device on load

jbloomAus · Nov 10, 2024 · 6e9350d · 6e9350d
1 parent 0c3a5b8
commit 6e9350d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sae_lens/load_model.py b/sae_lens/load_model.py
@@ -50,7 +50,7 @@ def load_model(
     elif model_class_name == "AutoModelForCausalLM":
         hf_model = AutoModelForCausalLM.from_pretrained(
             model_name, **model_from_pretrained_kwargs
-        )
+        ).to(device)
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         return HookedProxyLM(hf_model, tokenizer)