import torchvision during Predict warmup (#13313)
This commit is contained in:
parent
2684bcdc7d
commit
90f2a4226c
1 changed files with 2 additions and 0 deletions
|
|
@ -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…
Add table
Add a link
Reference in a new issue