coreml convert flatten to reshape, but npu does not support reshape

I have a model that uses ‘flatten’, and when I converted it to a Core ML model and ran it on Xcode with an iPhone XR, I noticed that ‘flatten’ was automatically converted to ‘reshape’. However, the NPU does not support ‘reshape’.

howerver, I got the Resnet50 model on apple models and performance it on XCode with the same iphone XR, I can see the 'flatten' operator which run on NPU.

On the other hand, when I used the following code to convert ResNet50 in PyTorch and ran it on Xcode Performance, the ‘flatten’ operation was converted to ‘reshape’, which then ran on the CPU.

? So I dont know how to keep 'flatten' operator when convert to ml model ?

coreml tool 7.1

iphone XR

ios 17.5.1

from torchvision import models
import coremltools as ct
import torch
import torch.nn as nn

network_name = "my_resnet50"
torch_model = models.resnet50(pretrained=True)
torch_model.eval()

width  = 224
height = 224

example_input = torch.rand(1, 3, height, width) 

traced_model = torch.jit.trace(torch_model,  (example_input))

model = ct.convert(
    traced_model,
    convert_to = "neuralnetwork", 
    inputs=[
            ct.TensorType(
                name            =  "data",
                shape           =  example_input.shape,
                dtype           =  np.float32
                )
        ],
    outputs = [
            ct.TensorType(
                name            =  "output",
                dtype           =  np.float32
            )
        ],
    compute_units               = ct.ComputeUnit.CPU_AND_NE,
    minimum_deployment_target   = ct.target.iOS14,
)
model.save("my_resnet.mlmodel")

ResNet50 on Resnet50.mlmodel

My Convertion of ResNet50