|
|
@ -620,6 +620,8 @@ class AutoBackend(nn.Module): |
|
|
|
Args: |
|
|
|
Args: |
|
|
|
imgsz (tuple): The shape of the dummy input tensor in the format (batch_size, channels, height, width) |
|
|
|
imgsz (tuple): The shape of the dummy input tensor in the format (batch_size, channels, height, width) |
|
|
|
""" |
|
|
|
""" |
|
|
|
|
|
|
|
import torchvision # noqa (import here so torchvision import time not recorded in postprocess time) |
|
|
|
|
|
|
|
|
|
|
|
warmup_types = self.pt, self.jit, self.onnx, self.engine, self.saved_model, self.pb, self.triton, self.nn_module |
|
|
|
warmup_types = self.pt, self.jit, self.onnx, self.engine, self.saved_model, self.pb, self.triton, self.nn_module |
|
|
|
if any(warmup_types) and (self.device.type != "cpu" or self.triton): |
|
|
|
if any(warmup_types) and (self.device.type != "cpu" or self.triton): |
|
|
|
im = torch.empty(*imgsz, dtype=torch.half if self.fp16 else torch.float, device=self.device) # input |
|
|
|
im = torch.empty(*imgsz, dtype=torch.half if self.fp16 else torch.float, device=self.device) # input |
|
|
|