henry000 commited on
Commit
7ad4a9d
Β·
2 Parent(s): 77181ee b74f67e

πŸ”€ [Merge] branch 'TEST' into DOCKER

Browse files
Files changed (1) hide show
  1. yolo/utils/model_utils.py +5 -0
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