import `torchvision` during Predict warmup (#13313)

pull/13315/head
Glenn Jocher 6 months ago committed by GitHub
parent 2684bcdc7d
commit 90f2a4226c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
  1. 2
      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

Loading…
Cancel
Save