From 15f6e57eaa6f003d12a8f403d2c6fe23610e76a3 Mon Sep 17 00:00:00 2001
From: Charles Goddard <chargoddard@gmail.com>
Date: Sun, 13 Aug 2023 13:19:48 -0700
Subject: [PATCH] Fix crash when running without CUDA

---
 src/axolotl/utils/bench.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/axolotl/utils/bench.py b/src/axolotl/utils/bench.py
index c39d07772..d19e81ecd 100644
--- a/src/axolotl/utils/bench.py
+++ b/src/axolotl/utils/bench.py
@@ -28,6 +28,9 @@ def gpu_memory_usage_smi(device=0):
 
 
 def log_gpu_memory_usage(log, msg, device):
+    if not torch.cuda.is_available():
+        return (0, 0, 0)
+
     usage, cache, misc = gpu_memory_usage_all(device)
     extras = []
     if cache > 0: