glenn-jocher commited on
Commit
d94b470
·
unverified ·
1 Parent(s): 8983324

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

Files changed (1) hide show
  1. 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