π [Merge] branch 'TEST' into DOCKER
Browse files
yolo/utils/model_utils.py
CHANGED
@@ -94,6 +94,11 @@ def get_device(device_spec: Union[str, int, List[int]]) -> torch.device:
|
|
94 |
device_spec = initialize_distributed()
|
95 |
if torch.cuda.is_available() and "cuda" in str(device_spec):
|
96 |
return torch.device(device_spec), ddp_flag
|
|
|
|
|
|
|
|
|
|
|
97 |
device = torch.device(device_spec)
|
98 |
return device, ddp_flag
|
99 |
|
|
|
94 |
device_spec = initialize_distributed()
|
95 |
if torch.cuda.is_available() and "cuda" in str(device_spec):
|
96 |
return torch.device(device_spec), ddp_flag
|
97 |
+
if not torch.cuda.is_available() and not torch.backends.mps.is_available():
|
98 |
+
if device_spec != "cpu":
|
99 |
+
logger.warning(f"β Device spec: {device_spec} not support, Choosing CPU instead")
|
100 |
+
return torch.device("cpu"), False
|
101 |
+
|
102 |
device = torch.device(device_spec)
|
103 |
return device, ddp_flag
|
104 |
|