Unable to compile PyTorch quantized models with torch-mlir due to CustomClass

What happened?

I'm trying to convert a quantized ResNet-50 model (using PyTorch's quantization-aware training) to MLIR format, but I'm encountering an error when torch-mlir attempts to import the model.

Environment

• python: 3.11.3
• pytorch: 2.7.1+cpu
• torch-mlir: 1983b6db9cdba2dfef4939a4fb521e9ac25dc08e

Steps to reproduce the issue

import torch
import torch.nn as nn
import torch.quantization
import torchvision
import torchvision.transforms as transforms
import torch_mlir.torchscript
import os

transform = 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]),
])
train_dataset = torchvision.datasets.FakeData(transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)

model = torchvision.models.quantization.resnet50(pretrained=True, quantize=False)

model.qconfig = torch.quantization.get_default_qat_qconfig("fbgemm")
torch.quantization.prepare_qat(model, inplace=True)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.train()

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)

 # QAT
for epoch in range(1): 
    for images, targets in train_loader:
        images, targets = images.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
    print(f"[Epoch {epoch+1}] QAT training complete")

model.eval()
quantized_model = torch.quantization.convert(model.eval(), inplace=False)

print(quantized_model)

scripted_model = torch.jit.script(quantized_model)
example_input = torch.randn(1, 3, 224, 224)

mlir_module = torch_mlir.torchscript.compile(
    scripted_model, 
    [example_input], 
    output_type="torch"
)

with open('resnet50_int8.mlir', 'w') as f:
    f.write(str(mlir_module))

log:

QuantizableResNet(
  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (relu): ReLU()
  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  (layer1): Sequential(
    (0): QuantizableBottleneck(
      (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU()
      (downsample): Sequential(
        (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
      (skip_add_relu): FloatFunctional(
        (activation_post_process): Identity()
      )
      (relu1): ReLU()
      (relu2): ReLU()
    )
...

Traceback (most recent call last):
  File "/workspace/qat_resnet50.py", line 57, in <module>
    mlir_module = torch_mlir.torchscript.compile(
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/path/to/torch-mlir/build/python_packages/torch_mlir/torch_mlir/torchscript.py", line 282, in compile
    raise Exception(
Exception: 
PyTorch TorchScript module -> torch-mlir Object Graph IR import failed with:
### Importer C++ Exception:
see diagnostics
### Importer Diagnostics:
error: unable to import Torch CustomClass type '0x5baca5da1580' to MLIR type

Do we have any approach to convert a quantized model to MLIR currently?

Reproduce the provided quantized ResNet-50 script and inspect torch_mlir/torchscript.py around line 282, where compilation reports the importer failure. Trace the Torch CustomClass diagnostic to determine whether quantized models are supported; done means documenting a verified conversion path or a clearly scoped unsupported case with a focused regression test if appropriate.