code: def init_fn(num_models):
      models = [MLPClassifier().to(DEVICE) for _ in range(num_models)]
      params, _ = stack_module_state(models)
      return params
  
file_path: pytorch/functorch/examples/ensembling/parallel_train.py
imports: import argparse
  import math
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.func import functional_call, grad_and_value, stack_module_state, vmap
label_desc: Use image node_id 4 for calling a global function with example usage: init_fn(num_models) and returns: params