Gracefully handle running with cuda torch on CPU only devices

This commit is contained in:
doctorpangloss 2024-10-10 10:42:22 -07:00
parent c34403b574
commit 5f26b76f59

View File

@ -1225,6 +1225,9 @@ def should_use_bf16(device=None, model_params=0, prioritize_performance=True, ma
props_major = min(torch.cuda.get_device_properties(torch.device(f"cuda:{i}")).major for i in range(torch.cuda.device_count()))
if props_major >= 8:
return True
except ValueError:
logging.warning("No CUDA devices were present, even though CUDA is available in this torch installation. This assumes the CPU device will be selected for computation")
return False
except AssertionError:
logging.warning("Torch was not compiled with CUDA support")
return False