diff --git a/models/common.py b/models/common.py index 049dfc0b9e0..8db00b9416a 100644 --- a/models/common.py +++ b/models/common.py @@ -688,7 +688,7 @@ def forward(self, im, augment=False, visualize=False): 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[x].data for x in sorted(self.output_names)] + y = [self.bindings[x].data for x in ["output", "462"]] elif self.coreml: # CoreML im = im.cpu().numpy() im = Image.fromarray((im[0] * 255).astype("uint8"))