Update hf.py (#284)

m0radt · web-flow · commit dac1da40276a · 2025-08-19T15:29:46.000-07:00
diff --git a/evalplus/provider/hf.py b/evalplus/provider/hf.py
@@ -23,6 +23,7 @@ def __init__(
     ):
         super().__init__(name=name, **kwargs)
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.device_map = device_map
 
         kwargs = {
             "device_map": device_map,
@@ -72,7 +73,7 @@ def codegen(
             )
         )
         input_tokens = self.tokenizer.encode(prompt, return_tensors="pt")
-        if device_map is None:
+        if self.device_map is None:
             input_tokens.to(self.device)
         kwargs = {}
         if do_sample:
