TRT `--half` fix autocast images to FP16 (#8435)
Browse files* TRT `--half` fix autocast images to FP16
Resolves bug raised in https://github.com/ultralytics/yolov5/issues/7822
* Update common.py
- models/common.py +3 -0
models/common.py
CHANGED
@@ -441,6 +441,9 @@ class DetectMultiBackend(nn.Module):
|
|
441 |
def forward(self, im, augment=False, visualize=False, val=False):
|
442 |
# YOLOv5 MultiBackend inference
|
443 |
b, ch, h, w = im.shape # batch, channel, height, width
|
|
|
|
|
|
|
444 |
if self.pt: # PyTorch
|
445 |
y = self.model(im, augment=augment, visualize=visualize)[0]
|
446 |
elif self.jit: # TorchScript
|
|
|
441 |
def forward(self, im, augment=False, visualize=False, val=False):
|
442 |
# YOLOv5 MultiBackend inference
|
443 |
b, ch, h, w = im.shape # batch, channel, height, width
|
444 |
+
if self.fp16 and im.dtype != torch.float16:
|
445 |
+
im = im.half() # to FP16
|
446 |
+
|
447 |
if self.pt: # PyTorch
|
448 |
y = self.model(im, augment=augment, visualize=visualize)[0]
|
449 |
elif self.jit: # TorchScript
|