Can you control the outputs when converting pytorch to CoreML?

Question

❓Question

I'm trying to convert a pytorch model to coreml. The model was based on yolov5.

Here is a netron view our our exported coreml model. Currently, the architecture has 3 outputs. You can see one of the outputs in the screenshot, number '740'.

However, we want a different output from coreml. We need to get the output before the reshapeStatic and transpose layers. So in this image you can see that we need the last convolution layer instead of 740.

Those reshapeStatic and transpose layers were added by the process which convert the net to coreML, they are not organic layers of yolov5.

1. Is there any way we can do the conversion to coreml differently in order to have more control over which layers are output? For example, can we have control over the output layers in the sample code below:

model = ct.convert(
    traced_model,
    inputs=[ct.ImageType(name="input_1", shape=example_input.shape)], 
    classifier_config = ct.ClassifierConfig(class_labels) )

2. Alternatively, is there a way where we can choose at runtime which values to pull out of the coreml model? For example, is there a way to specify in the code below which layers we want to output?

img = load_image(img_path, resize_to=(img_size, img_size))
//can we specify here which layers to output?
coreml_out_dict = model.predict({'image': img})

Thanks!