Spaces:
Sleeping
Sleeping
💚 [Fix] MPS CI bugs, enable using MPS
Browse files
yolo/utils/model_utils.py
CHANGED
|
@@ -94,7 +94,7 @@ 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 |
-
if not torch.cuda.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
|
|
|
|
| 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():
|
| 98 |
if device_spec != "cpu":
|
| 99 |
logger.warning(f"❎ Device spec: {device_spec} not support, Choosing CPU instead")
|
| 100 |
return torch.device("cpu"), False
|