Skip to content

Commit c1249a4

Browse files
authored
*.torchscript inference self.jit fix (#6007)
1 parent 628817d commit c1249a4

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

models/common.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -379,7 +379,7 @@ def wrap_frozen_graph(gd, inputs, outputs):
379379
def forward(self, im, augment=False, visualize=False, val=False):
380380
# YOLOv5 MultiBackend inference
381381
b, ch, h, w = im.shape # batch, channel, height, width
382-
if self.pt: # PyTorch
382+
if self.pt or self.jit: # PyTorch
383383
y = self.model(im) if self.jit else self.model(im, augment=augment, visualize=visualize)
384384
return y if val else y[0]
385385
elif self.coreml: # CoreML

0 commit comments

Comments
 (0)