code: def step6():
      parallel_train_step_fn = vmap(
          train_step_fn, in_dims=(0, None, None)
      )
      batched_weights = init_fn(num_models=2)
      for i in range(2000):
          loss, batched_weights = parallel_train_step_fn(
              batched_weights, points, labels
          )
          if i % 200 == 0:
              print(loss)
  
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 5 for calling a global function with example usage: step6() without return types