Error when trying to emit MLIR for Training MNIST kernel

Hi, I m trying to generate mlir for the training MNIST kernel using the below script
`
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.optim.lr_scheduler import StepLR
import torch_mlir
from functorch import make_fx
from torch.nn.utils import stateless
from torch._functorch.compile_utils import strip_overloads
from torch._decomp import get_decompositions

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)
        self.pool = nn.MaxPool2d(kernel_size=2)
        self.relu = nn.ReLU()
        self.logsoftmax = nn.LogSoftmax(dim=1)
        self.flatten = nn.Flatten(1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.relu(x)
        x = self.pool(x)
        x = self.dropout1(x)
        x = self.flatten(x)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = self.logsoftmax(x)
        return output

mod = Net()

def forward(params, buffers, args):
    params_and_buffers = {**params, **buffers}
    optimizer = optim.Adadelta(mod.parameters(), lr=1.0)
    optimizer.zero_grad()
    res = stateless.functional_call(mod, params_and_buffers, args,
                               {})
    loss = F.nll_loss(res, target)
    loss.backward()
    optimizer.step()
    return params, buffers

def get_sorted_params(named_params):
    return [i[1] for i in sorted(named_params.items())]

transform=transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
    ])
dataset1 = datasets.MNIST('../data', train=True, download=True,
                    transform=transform)
dataset2 = datasets.MNIST('../data', train=False,
                    transform=transform)
train_loader = torch.utils.data.DataLoader(dataset1)#,**train_kwargs)
test_loader = torch.utils.data.DataLoader(dataset2)#, **test_kwargs)

data = None
for batch_idx, (data, target) in enumerate(train_loader):
    data, target = data.to("cpu"), target.to("cpu")
    break

arg = []
for name, param in mod.named_parameters():
    arg.append(param.data)
arg.append(data)
fx_graph = make_fx(forward)(dict(mod.named_parameters()),
                        dict(mod.named_buffers()), data)
fx_graph.graph.set_codegen(torch.fx.graph.CodeGen())
fx_graph.recompile()
sinput = strip_overloads(fx_graph)

ts_graph = torch.jit.script(fx_graph)

linalg_on_tensors_mlir = torch_mlir.compile(
    ts_graph,
    arg,
    output_type=torch_mlir.OutputType.LINALG_ON_TENSORS, use_tracing=True)
print(linalg_on_tensors_mlir)

`

But I m getting the below error:
python exception: Failure while executing pass pipeline:
error: unknown: unsupported by backend contract: tensor with unknown rank
note: unknown: see current operation: %25 = "torch.tensor_static_info_cast"(%arg0) : (!torch.vtensor<[32,1,3,3],f32>) -> !torch.vtensor
note: unknown: this is likely due to a missing transfer function in abstract_interp_lib_gen.py

In the print after all, I noticed that "torch.tenosr_static_info_cast" op is getting created for each function argument during the AdjustCallingConventions pass.

Start with the provided MNIST training script and the torch_mlir.compile call using LINALG_ON_TENSORS. Trace the AdjustCallingConventions pass and the tensor_static_info_cast operation, then inspect abstract_interp_lib_gen.py for the missing transfer function noted by the diagnostic. Done means the script no longer produces an unknown-rank backend-contract error.