Fix TensorRT --dynamic excess outputs bug

Potential fix for #8790
ultralytics · Aug 4, 2022 · 9d9c2c6 · 9d9c2c6
1 parent 84e7748
commit 9d9c2c6
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/models/common.py b/models/common.py
@@ -480,6 +480,8 @@ def forward(self, im, augment=False, visualize=False, val=False):
             self.binding_addrs['images'] = int(im.data_ptr())
             self.context.execute_v2(list(self.binding_addrs.values()))
             y = self.bindings['output'].data
+            if len(y) > b:
+                y = y[:b]  # trim excess --dynamic outputs
         elif self.coreml:  # CoreML
             im = im.permute(0, 2, 3, 1).cpu().numpy()  # torch BCHW to numpy BHWC shape(1,320,192,3)
             im = Image.fromarray((im[0] * 255).astype('uint8'))