use torch.cuda.current_device() instead of local_rank

axolotl-ai-cloud · Sep 15, 2023 · 56da204 · 56da204
1 parent 9218ebe
commit 56da204
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/axolotl/utils/config.py b/src/axolotl/utils/config.py
@@ -29,7 +29,7 @@ def get_device():
         cfg.device_map = "auto"
     else:
         if cfg.device.startswith("cuda"):
-            cfg.device_map = {"": cfg.local_rank}
+            cfg.device_map = {"": torch.cuda.current_device()}
         else:
             cfg.device_map = {"": cfg.device}