torch.func.stack_module_state
torch.func.stack_module_state(models) -> params, buffers
Prepares a list of torch.nn.Modules for ensembling with "vmap()".
Given a list of "M" "nn.Modules" of the same class, returns two
   dictionaries that stack all of their parameters and buffers
   together, indexed by name. The stacked parameters are optimizable
   (i.e. they are new leaf nodes in the autograd history that are
   unrelated to the original parameters and can be passed directly to
   an optimizer).
Here's an example of how to ensemble over a very simple model:
  num_models = 5
  batch_size = 64
  in_features, out_features = 3, 3
  models = [torch.nn.Linear(in_features, out_features) for i in range(num_models)]
  data = torch.randn(batch_size, 3)

  def wrapper(params, buffers, data):
      return torch.func.functional_call(model[0], (params, buffers), data)

  params, buffers = stack_module_state(models)
