Single-line TRT dynamic assertion (#8871)
parent
38a6eb6e99
commit
731a2f8c1f
|
@ -477,8 +477,7 @@ class DetectMultiBackend(nn.Module):
|
|||
self.bindings['images'] = self.bindings['images']._replace(shape=im.shape)
|
||||
self.bindings['output'].data.resize_(tuple(self.context.get_binding_shape(i_out)))
|
||||
s = self.bindings['images'].shape
|
||||
assert im.shape == s, f"image shape {im.shape} " + \
|
||||
f"exceeds model max shape {s}" if self.dynamic else f"does not match model shape {s}"
|
||||
assert im.shape == s, f"input size {im.shape} {'>' if self.dynamic else 'not equal to'} max model size {s}"
|
||||
self.binding_addrs['images'] = int(im.data_ptr())
|
||||
self.context.execute_v2(list(self.binding_addrs.values()))
|
||||
y = self.bindings['output'].data
|
||||
|
|
Loading…
Reference in New Issue