chargoddard
commited on
Commit
·
15f6e57
1
Parent(s):
729c299
Fix crash when running without CUDA
Browse files
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:
|