Move model to GPU:0 for SqueezeLLM

danielz02 · Mar 26, 2024 · 55a5407 · 55a5407
1 parent a67724f
commit 55a5407
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/helm/proxy/clients/huggingface_client.py b/src/helm/proxy/clients/huggingface_client.py
@@ -111,7 +111,7 @@ def __init__(self, model_config: HuggingFaceModelConfig):
 
                     model = LlamaForCausalLM.from_pretrained(model_name, torch_dtype="auto")
 
-                model = model.eval()
+                model = model.to(self.device).eval()
                 layers = find_layers(model)
 
                 state_dict = torch.load(quantization_config.quant_file)