ntsc207 commited on
Commit
a4e7334
·
verified ·
1 Parent(s): 898ef8d

Update strong_sort/reid_multibackend.py

Browse files
Files changed (1) hide show
  1. strong_sort/reid_multibackend.py +2 -0
strong_sort/reid_multibackend.py CHANGED
@@ -124,6 +124,8 @@ class ReIDDetectMultiBackend(nn.Module):
124
  return pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs
125
 
126
  def warmup(self, imgsz=(1, 256, 128, 3)):
 
 
127
  # Warmup model by running inference once
128
  warmup_types = self.pt, self.jit, self.onnx, self.engine, self.saved_model, self.pb
129
  if any(warmup_types) and self.device.type != 'cpu':
 
124
  return pt, jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs
125
 
126
  def warmup(self, imgsz=(1, 256, 128, 3)):
127
+ if isinstance(self.device, str):
128
+ self.device = torch.device(self.device if torch.cuda.is_available() else 'cpu')
129
  # Warmup model by running inference once
130
  warmup_types = self.pt, self.jit, self.onnx, self.engine, self.saved_model, self.pb
131
  if any(warmup_types) and self.device.type != 'cpu':