glenn-jocher commited on
Commit
c1249a4
1 Parent(s): 628817d

*.torchscript inference `self.jit` fix (#6007)

Browse files
Files changed (1) hide show
  1. models/common.py +1 -1
models/common.py CHANGED
@@ -379,7 +379,7 @@ class DetectMultiBackend(nn.Module):
379
  def forward(self, im, augment=False, visualize=False, val=False):
380
  # YOLOv5 MultiBackend inference
381
  b, ch, h, w = im.shape # batch, channel, height, width
382
- if self.pt: # PyTorch
383
  y = self.model(im) if self.jit else self.model(im, augment=augment, visualize=visualize)
384
  return y if val else y[0]
385
  elif self.coreml: # CoreML
 
379
  def forward(self, im, augment=False, visualize=False, val=False):
380
  # YOLOv5 MultiBackend inference
381
  b, ch, h, w = im.shape # batch, channel, height, width
382
+ if self.pt or self.jit: # PyTorch
383
  y = self.model(im) if self.jit else self.model(im, augment=augment, visualize=visualize)
384
  return y if val else y[0]
385
  elif self.coreml: # CoreML