code: def run_test_collect_shards(
      rank, world_size, reference_rank, tempfile_name
  ):
      dist_init(rank, world_size, tempfile_name)
      device = (
          torch.device(rank)
          if torch.cuda.device_count() > 1
          else DEVICE
      )
      torch.cuda.set_device(rank)
  
      # Run a dummy step so that the optimizer state dict exists
      batch, input_width, hidden, target_width = 3, 3, 3, 5
      target = torch.rand((batch, target_width), device=device)
      inputs = torch.rand((batch, input_width), device=device)
  
      model = torch.nn.Sequential(
          torch.nn.Linear(input_width, hidden),
          torch.nn.Linear(hidden, target_width),
      )
      model.to(device)
  
      loss_fn = torch.nn.L1Loss()
      loss_fn.to(device)
  
      # With SGD, Momentum is required to get a state to shard
      optimizer = optim.OSS(model.parameters(), lr=0.1, momentum=0.99)
  
      def closure():
          optimizer.zero_grad()
          output = model(inputs)
          loss = loss_fn(output, target)
          loss.backward()
          return loss
  
      _ = optimizer.step(closure=closure)
  
      # Update the optimizer state on the reference rank
      optimizer.consolidate_state_dict(recipient_rank=reference_rank)
  
      # Fetch the state on the reference rank
      # - check that it has the correct size
      # - load it again
      if rank == reference_rank:
          optimizer_state_dict = optimizer.state_dict()
          assert len(optimizer_state_dict["state"]) == len(
              list(model.parameters())
          )
      else:
          optimizer_state_dict = {}
  
      # distribute to the other ranks
      optimizer_state_dict = sync_object_ranks(
          optimizer_state_dict, reference_rank, device
      )
  
      # Load the optimizer state dict
      optimizer.load_state_dict(optimizer_state_dict)
  
      # Check that the states are not None, but {}
      for state in optimizer.state.values():
          for _, _ in state.items():
              pass
  
      # Test the state dict materialization on all ranks
      _ = optimizer.step(closure=closure)
      optimizer_state_dict = optimizer.state_dict(
          all_ranks=True
      )  # one per rank
      optimizer.load_state_dict(optimizer_state_dict)
      _ = optimizer.step(closure=closure)
      check_same_models_across_ranks(
          model,
          dist.group.WORLD,
          params_should_be_equal=True,
          check_broadcast_buffers=False,
      )
  
      # Check that if the model is moved to cpu, the optimizer consolidation still works
      model.cpu()
      optimizer = optim.OSS(model.parameters(), lr=0.1, momentum=0.99)
      optimizer.consolidate_state_dict(recipient_rank=reference_rank)
  
      dist.destroy_process_group()
  
imports: import copy
  from math import inf
  import tempfile
  from typing import Any, Dict, Type, cast
  import unittest
  import numpy
  import pytest
  import torch
  import torch.distributed
  import torch.multiprocessing
  from torch.nn.parallel import DistributedDataParallel
  from fairscale.fair_dev.testing.testing import check_same_model_params, check_same_models_across_ranks, skip_if_no_cuda, skip_if_py39_no_cuda, skip_if_single_gpu
  from fairscale.internal import torch_version
  import fairscale.optim
label_desc: Use image node_id 15 for calling a global function with example usage: run_test_collect_shards(rank, world_size, reference_rank, tempfile_name) and returns: loss