code: def _train(model, in_data, steps_per_iter=1):
      optim = SGD(model.parameters(), lr=0.1)
      for _ in range(3):
          # Simulate multiple ranks.
          for _ in range(steps_per_iter):
              out = model(in_data)
              out.backward()
          # Simulate gradient means between ranks.
          if steps_per_iter > 1:
              with torch.no_grad():
                  for p in model.parameters():
                      p.grad /= steps_per_iter
          with torch.no_grad():
              for p in model.parameters():
                  assert (
                      not torch.isinf(p.grad).any()
                      and not torch.isnan(p.grad).any()
                  )
          optim.step()
          model.zero_grad(set_to_none=True)
  
file_path: fairscale/tests/nn/data_parallel/test_fsdp_shared_weights_mevo.py
imports: from copy import deepcopy
  import pytest
  import torch
  from torch import nn
  import torch.multiprocessing
  from torch.optim import SGD
  from fairscale.experimental.nn import MEVO
  from fairscale.fair_dev.testing.testing import dist_init, in_circle_ci, objects_are_equal, skip_if_single_gpu, teardown, temp_files_ctx
  from fairscale.nn.data_parallel import FullyShardedDataParallel
label_desc: Use image node_id 5 for calling a global function with example usage: _train(model, in_data, steps_per_iter) without return types