henry000 commited on
Commit
3b31306
·
1 Parent(s): 7cbff17

🚑️ [Fix] get_device typo, not return cpu

Browse files
Files changed (1) hide show
  1. yolo/utils/model_utils.py +2 -1
yolo/utils/model_utils.py CHANGED
@@ -94,7 +94,8 @@ 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
- return torch.device("cpu"), False
 
98
 
99
 
100
  class PostProccess:
 
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
 
100
 
101
  class PostProccess: