skip gpu mem logging if cpu too

axolotl-ai-cloud · Sep 21, 2023 · b1f1f44 · b1f1f44
1 parent 972bd42
commit b1f1f44
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/axolotl/utils/bench.py b/src/axolotl/utils/bench.py
@@ -19,7 +19,7 @@ def actual_decorator(func):
         def wrapper(*args, **kwargs):
             device = kwargs.get("device", args[0] if args else None)
 
-            if not torch.cuda.is_available() or device == "auto":
+            if not torch.cuda.is_available() or device == "auto" or device == "cpu":
                 return default_value
 
             return func(*args, **kwargs)