also wrap log_gpu_memory_usage

axolotl-ai-cloud · Sep 21, 2023 · 8efef9a · 8efef9a
1 parent b1f1f44
commit 8efef9a
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/src/axolotl/utils/bench.py b/src/axolotl/utils/bench.py
@@ -57,10 +57,8 @@ def gpu_memory_usage_smi(device=0):
         return 0.0
 
 
+@check_cuda_device((0, 0, 0))
 def log_gpu_memory_usage(log, msg, device):
-    if not torch.cuda.is_available() or device == "auto":
-        return (0, 0, 0)
-
     usage, cache, misc = gpu_memory_usage_all(device)
     extras = []
     if cache > 0: