TRT --half fix autocast images to FP16 (#8435)

* TRT `--half` fix autocast images to FP16

Resolves bug raised in https://github.com/ultralytics/yolov5/issues/7822

* Update common.py
This commit is contained in:
Glenn Jocher authored and GitHub committed 2022-07-01 15:41:14 +02:00
1 parent 898332433a
commit d94b4705a6
1 file changed
+3
+3
View File
@@ -441,6 +441,9 @@ class DetectMultiBackend(nn.Module):
def forward(self, im, augment=False, visualize=False, val=False):
# YOLOv5 MultiBackend inference
b, ch, h, w = im.shape # batch, channel, height, width
if self.fp16 and im.dtype != torch.float16:
im = im.half() # to FP16
if self.pt: # PyTorch
y = self.model(im, augment=augment, visualize=visualize)[0]
elif self.jit: # TorchScript