diff --git a/models/common.py b/models/common.py index c898d94a921a..f93ef8736d7b 100644 --- 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'))