From 6c468e358e597f6926398dbdc518c48c8e32d380 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 4 Aug 2022 23:33:04 +0200 Subject: [PATCH] Single-line TRT dynamic assertion --- models/common.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/models/common.py b/models/common.py index cfa688ba940b..a1269c5f3372 100644 --- a/models/common.py +++ b/models/common.py @@ -477,8 +477,7 @@ def forward(self, im, augment=False, visualize=False, val=False): 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