Fix Trainer when model is loaded on a different GPU (huggingface#23792)

Lucinao540 · May 31, 2023 · 68d53bc · 68d53bc
1 parent 0963a25
commit 68d53bc
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
@@ -405,12 +405,12 @@ def __init__(
         else:
             self.is_model_parallel = False
 
-        if (
-            getattr(model, "hf_device_map", None) is not None
-            and len([device for device in set(model.hf_device_map.values()) if device not in ["cpu", "disk"]]) > 1
-            and not self.is_model_parallel
-        ):
-            self.is_model_parallel = True
+        if getattr(model, "hf_device_map", None) is not None:
+            devices = [device for device in set(model.hf_device_map.values()) if device not in ["cpu", "disk"]]
+            if len(devices) > 1:
+                self.is_model_parallel = True
+            else:
+                self.is_model_parallel = self.args.device != torch.device(devices[0])
 
             # warn users
             logger.info(