Understanding shape of example arguments to forward method in PyTorch

I am trying to compile my PyTorch model into mlir using torch-mlir. The compile method from torch-mlir requires three arguments: 1) model to convert 2) example arguments to use when inferring the shapes of the arguments to the forward method of the model 3) desired output type.

My question is with the second input. I am not sure what the shape is of the arguments to the forward method. The model I am using to practice with is the Iris dataset. The full code is here.

The model and forward method are:

class Model(nn.Module):
    def __init__(self, input_dim):
        super(Model, self).__init__()
        self.layer1 = nn.Linear(input_dim, 50)
        self.layer2 = nn.Linear(50, 3)
        
    def forward(self, x):
        print(x.shape)
        x = F.relu(self.layer1(x))
        x = F.softmax(self.layer2(x), dim=1)
        return x

From the repository you provided:

example_args: A list of example arguments to use when inferring the shapes of the arguments to forward method of the model. A single tensor is treated as a list of a single tensor. A TensorPlaceholder object is also allowed in the place of any Tensor. For models with multiple methods, an ExampleArgs object can be passed.

This means that if x has a shape of (sample_size, input_dims), the argument for example_args should be (number_of_examples, sample_size, input_dims).