fix handling for device

axolotl-ai-cloud · Sep 7, 2023 · b67fe7c · b67fe7c
1 parent 898c70a
commit b67fe7c
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/axolotl/utils/bench.py b/src/axolotl/utils/bench.py
@@ -28,7 +28,7 @@ def gpu_memory_usage_smi(device=0):
 
 
 def log_gpu_memory_usage(log, msg, device):
-    if not torch.cuda.is_available():
+    if not torch.cuda.is_available() or device == "auto":
         return (0, 0, 0)
 
     usage, cache, misc = gpu_memory_usage_all(device)

diff --git a/src/axolotl/utils/config.py b/src/axolotl/utils/config.py
@@ -24,10 +24,10 @@ def get_device():
         except Exception:  # pylint: disable=broad-exception-caught
             return "cpu"
 
+    cfg.device = get_device()
     if cfg.world_size == 1:
         cfg.device_map = "auto"
     else:
-        cfg.device = get_device()
         if cfg.device.startswith("cuda"):
             cfg.device_map = {"": cfg.local_rank}
         else: