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:
|