From 90f2a4226c6bfd87f2167d559aa6f23d9b4f49c7 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 3 Jun 2024 00:54:10 +0200 Subject: [PATCH] import `torchvision` during Predict warmup (#13313) --- ultralytics/nn/autobackend.py | 2 ++ 1 file changed, 2 insertions(+) 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