Error compiling MobilenetV2

I'm trying to compile MobilenetV2 with torch_mlir, but I get the following error

Traceback (most recent call last):
  File ".........../mobilnet/main.py", line 44, in <module>
    module = torch_mlir.compile(model, [input_batch] ,output_type=torch_mlir.OutputType.TOSA, use_tracing=True)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/marco/.local/lib/python3.11/site-packages/torch_mlir/__init__.py", line 444, in compile
    raise Exception(f"""
Exception:
PyTorch TorchScript module -> torch-mlir Object Graph IR import failed with:
### Importer C++ Exception:
bad_weak_ptr
### Importer Diagnostics:

Code is:

import torch
import torch_mlir
import torchvision.models as models
from torchvision import transforms
from PIL import Image
#from torchvision import transforms
import urllib

class Wrapper(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        # Reset seed to make model deterministic.
        torch.manual_seed(0)
        self.mobilenet = models.mobilenet_v2()
        self.train(False)

    def forward(self, img):
        return self.mobilenet.forward(img)

#model = torch.hub.load('pytorch/vision:v0.10.0', 'mobilenet_v3', pretrained=True)
model = Wrapper()
model.eval()

# Download an example image from the pytorch website
url, filename = ("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")
try: urllib.URLopener().retrieve(url, filename)
except: urllib.request.urlretrieve(url, filename)
#
input_image = Image.open(filename)
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model

# DO NOT TOUCH
out_tosa_mlir_path = "./mobilenet.mlir"

#batch, tensor, input_image?
module = torch_mlir.compile(model, [input_batch] ,output_type=torch_mlir.OutputType.TOSA, use_tracing=True)
with open(out_tosa_mlir_path, "w", encoding="utf-8") as outf:
    outf.write(str(module))

Torch and Torch-MLIR versions are: torch-2.3.0.dev20240122+cpu-cp311-cp311-linux_x86_64.whl torch_mlir-20240127.1096-cp311-cp311-linux_x86_64.whl

llvm / torch-mlir

Error compiling MobilenetV2 #3533