diff --git a/ultralytics/nn/autobackend.py b/ultralytics/nn/autobackend.py index bd07d1d4..d2727fe2 100644 --- a/ultralytics/nn/autobackend.py +++ b/ultralytics/nn/autobackend.py @@ -620,6 +620,8 @@ class AutoBackend(nn.Module): Args: 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 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