Spaces:
Running
on
Zero
Running
on
Zero
Update strong_sort/reid_multibackend.py
Browse files
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':
|