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:
1 parent
898332433a
commit
d94b4705a6
1 file changed
+3
@@ -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
|
||||
|
||||
Reference in new issue
Block a user