code: def step4():
      global weights
      for i in range(2000):
          loss, weights = train_step_fn(
              dict(model.named_parameters()), points, labels
          )
          if i % 100 == 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 3 for calling a global function with example usage: step4() without return types