Ankle/yolov5_birdproject
1
0 Fork 0
Code Issues Pull Requests Actions 6 Packages Projects Releases Wiki Activity

Single-line TRT dynamic assertion (#8871)

Browse Source
This commit is contained in:
Glenn Jocher authored and GitHub committed 2022-08-04 23:34:15 +02:00
1 parent 38a6eb6e99
commit 731a2f8c1f
1 file changed
+1 -2
Download Patch File Download Diff File Expand all files Collapse all files
models/common.py
+1 -2
View File
@@ -477,8 +477,7 @@ class DetectMultiBackend(nn.Module):
self.bindings['images'] = self.bindings['images']._replace(shape=im.shape) self.bindings['images'] = self.bindings['images']._replace(shape=im.shape)
self.bindings['output'].data.resize_(tuple(self.context.get_binding_shape(i_out))) self.bindings['output'].data.resize_(tuple(self.context.get_binding_shape(i_out)))
s = self.bindings['images'].shape s = self.bindings['images'].shape
assert im.shape == s, f"image shape {im.shape} " + \ assert im.shape == s, f"input size {im.shape} {'>' if self.dynamic else 'not equal to'} max model size {s}"
f"exceeds model max shape {s}" if self.dynamic else f"does not match model shape {s}"
self.binding_addrs['images'] = int(im.data_ptr()) self.binding_addrs['images'] = int(im.data_ptr())
self.context.execute_v2(list(self.binding_addrs.values())) self.context.execute_v2(list(self.binding_addrs.values()))
y = self.bindings['output'].data y = self.bindings['output'].data