zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

TensorRT PyTorch Hub inference fix (#7560) 

Solution proposed in https://github.com/ultralytics/yolov5/issues/7128 to TRT PyTorch Hub CUDA illegal memory errors.
This commit is contained in:
Glenn Jocher 2022-04-24 12:45:56 -07:00 committed by GitHub
parent c16671fc74
commit 950a85d9f6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
models/common.py
View File

@ -531,7 +531,7 @@ class AutoShape(nn.Module):
# multiple: = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...] # list of images
t = [time_sync()]
p = next(self.model.parameters()) if self.pt else torch.zeros(1) # for device and type
p = next(self.model.parameters()) if self.pt else torch.zeros(1, device=self.model.device) # for device, type
autocast = self.amp and (p.device.type != 'cpu') # Automatic Mixed Precision (AMP) inference
if isinstance(imgs, torch.Tensor): # torch
with amp.autocast(autocast):