    def from_input_list_with_batch_size(
        cls, inputs: List[torch.Tensor], max_batch_size: int, batch_dim: int = 0
    ) -> List["TensorSpec"]:
        """
        Most of the recommendation models will work fine using this function.

        We make an assumption that inferred lowerable subgraph inputs will have
        a single batch dimension with the same max batch size.
        """
        result: List[TensorSpec] = []

        bs_dim = cls.find_batch_size_dim(inputs)
        for index, t in enumerate(inputs):
            shape: List[IntVar] = []
            for i, d in enumerate(t.shape):
                if i == bs_dim[index]:
                    shape.append(IntVar([1, max_batch_size], "batch_size"))
                else:
                    shape.append(IntImm(d))
            result.append(TensorSpec(shape, t.dtype))

        return result
Most of the recommendation models will work fine using this function.

We make an assumption that inferred lowerable subgraph inputs will have
a single batch dimension with the same max batch size.
fx2ait/fx2ait/tensor_spec.py