I can't convert MobileVitV2 to CoreML

[zzmicer]

I've trained mobilevitv2_200_384_in22ft1k on custom dataset. Now trying to convert it to CoreML via coremltools. Get an error Encountered an error while compiling a neural network model: in operation of type reshape: Rank of the shape parameter must be between 0 and 5 (inclusive) in reshape

Versions
torch: 1.12.1
python: 3.8.8
coremltools: 5.2.0
To Reproduce
Steps to reproduce the behavior:

1. Load model weights

model = (
    torch.load("mobile_vit_last.pth",map_location=torch.device('cpu'))
)
model.eval()

2. Gather Jit trace

import timm
import torch
from torchvision.transforms import functional as F

from PIL import Image

img = Image.open("sample_image.jpg").convert('RGB')
img = F.resize(
            img=img, size=384, interpolation=F.InterpolationMode.BILINEAR
        )
input_tensor = F.pil_to_tensor(img).float()
input_tensor.div_(255.0)
input_tensor = input_tensor.unsqueeze(0)

with torch.no_grad():
    jit_model = torch.jit.trace(model, input_tensor)

3. Convert to CoreML

 coreml_model = ct.convert(
            model=jit_model,
            inputs=[
                ct.ImageType(name="input", shape=input_tensor.shape, scale=1.0 / 255.0)
            ],
            convert_to="mlprogram",
        )

4. Try on sample image

img = Image.open("sample_image.jpg").convert('RGB').resize((384,384))
# Make a prediction using Core ML
out_dict = coreml_model.predict(tensor)

Expected behavior

[rwightman]

@zzmicer unforunately you're on your own with this one, I don't have the bandwidth to work with CoreML and I do know it has more limitations right now than say ONNX / tensorrt. The original code might have some extra workarounds / differences for coreml as it's by Apple. This model does export to ONNX.

I'll move to discussions in case anyone else has any input. I don't consider it a bug at this point but a CoreML limitation.

[zzmicer]

@rwightman thanks for a fast reply!
I'll try to deploy that nn via the TorchScript for the moment.