[TorchDynamo] Too many function args in a fx graph

Hi all, I’m working on compiling a model using Torchdynamo + Torch-MLIR, and running it via IREE. The issue I’m facing is that the fx graph that is getting generated contains nearly 1000 function args (all tensors). Because of such a large number of function args, it results in a segmentation fault during IREE runtime. The reason is that the system heap gets exhausted when allocating such a large number of tensors (which are large in size). FYI, these tensors are the weights of the model. I have two questions here:
1.) Can we make any changes to the fx graph such that these function args are replaced by a single arg which would be a list of tensors and all the usages of these function args with the list accesses?
2.) Is there any other way we can handle this type of instance via torchdynamo?

cc @SherlockNoMad who might know

You are getting a FX graph produced by AOTAutograd, which lifts parameters/buffers as graph inputs. You may ask why we wanna do this? This is because we want to produce a stateless graph, which will be easier to work with for most backends.

As a work around, you may want to write a pass that fold inputs back into graph as its internal state.