chargoddard commited on
Commit
15f6e57
·
1 Parent(s): 729c299

Fix crash when running without CUDA

Browse files
Files changed (1) hide show
  1. src/axolotl/utils/bench.py +3 -0
src/axolotl/utils/bench.py CHANGED
@@ -28,6 +28,9 @@ def gpu_memory_usage_smi(device=0):
28
 
29
 
30
  def log_gpu_memory_usage(log, msg, device):
 
 
 
31
  usage, cache, misc = gpu_memory_usage_all(device)
32
  extras = []
33
  if cache > 0:
 
28
 
29
 
30
  def log_gpu_memory_usage(log, msg, device):
31
+ if not torch.cuda.is_available():
32
+ return (0, 0, 0)
33
+
34
  usage, cache, misc = gpu_memory_usage_all(device)
35
  extras = []
36
  if cache > 0: