fix device_map (#505)

EvolvingLMMs-Lab · Jan 23, 2025 · 9c4929f · 9c4929f
1 parent b1fbf55
commit 9c4929f
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/lmms_eval/models/internvl2.py b/lmms_eval/models/internvl2.py
@@ -204,8 +204,8 @@ def __init__(
             self._device = torch.device(f"cuda:{accelerator.local_process_index}")
             self.device_map = f"cuda:{accelerator.local_process_index}"
 
-        self._model = AutoModel.from_pretrained(self.path, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, trust_remote_code=True, device_map=device_map).eval()
-        self._tokenizer = AutoTokenizer.from_pretrained(self.path, trust_remote_code=True, device_map=device_map)
+        self._model = AutoModel.from_pretrained(self.path, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, trust_remote_code=True, device_map=self.device_map).eval()
+        self._tokenizer = AutoTokenizer.from_pretrained(self.path, trust_remote_code=True, device_map=self.device_map)
 
         if accelerator.num_processes > 1:
             assert accelerator.distributed_type in [DistributedType.FSDP, DistributedType.MULTI_GPU, DistributedType.DEEPSPEED], "Unsupported distributed type provided. Only DDP and FSDP are supported."

diff --git a/lmms_eval/models/qwen2_audio.py b/lmms_eval/models/qwen2_audio.py
@@ -58,7 +58,7 @@ def __init__(
         self._model = Qwen2AudioForConditionalGeneration.from_pretrained(
             pretrained,
             torch_dtype="auto",
-            device_map=device_map,
+            device_map=self.device_map,
         ).eval()
 
         self.processor = AutoProcessor.from_pretrained(pretrained)