code: def test_sequential(oss):
      """Test adascale with DDP + OSS with a sequential model"""
      world_size = 2
      temp_file_name = tempfile.mkstemp()[1]
  
      # Run multiple iterations, check the gain for both oss and non-oss cases.
      #
      # The inputs are picked arbitrarily. I used vectors that are orthogonal.
      #
      # The gain and mean_weight values are recorded from my testing and used here
      # to ensure their value is unchanged from commit to commit unless we can
      # explain why.
      test_case = {
          "inputs": [[[1.0, 0], [0, 1.0]], [[0, 1.0], [1.0, 0]]],
          "expected_gain": 1.0335265132125744,
          "expected_mean_weight": 52.92657661437988,
      }
  
      if oss == "oss-wrapper":
          # When OSS wraps AdaScale, the training is numerically different
          # and it exists only to enable future research. So we don't check
          # the gain (OSS doesn't have a gain() function, different rank's
          # gains are different). We just ensure the mean_weight is expected.
          del test_case["expected_gain"]
          test_case["expected_mean_weight"] = 94.93386840820312
  
      # The model.
      model = Sequential(
          Linear(2, 3, bias=False),
          Linear(3, 4, bias=False),
          Linear(4, 5, bias=False),
          Linear(5, 6, bias=False),
      )
  
      # Weights need to be fixed for deterministic gain.
      model[0].weight.data.copy_(
          Tensor(range(6)).reshape(3, 2) / mean(range(6))
      )
      model[1].weight.data.copy_(
          Tensor(range(12)).reshape(4, 3) / mean(range(12))
      )
      model[2].weight.data.copy_(
          Tensor(range(20)).reshape(5, 4) / mean(range(20))
      )
      model[3].weight.data.copy_(
          Tensor(range(30)).reshape(6, 5) / mean(range(30))
      )
  
      mp.spawn(
          _test_basic_func,
          args=(world_size, temp_file_name, test_case, oss, model),
          nprocs=world_size,
          join=True,
      )
  
file_path: fairscale/tests/optim/test_oss_adascale.py
imports: from statistics import mean
  import tempfile
  import numpy
  import pytest
  import torch
  from torch import Tensor
  import torch.distributed
  import torch.multiprocessing
  from torch.nn import Linear, Sequential
  from torch.nn.parallel import DistributedDataParallel
  from torch.optim import SGD
  from fairscale.fair_dev.testing.golden_testing_data import adascale_test_data
  from fairscale.fair_dev.testing.testing import skip_if_single_gpu
  from fairscale.optim import OSS, AdaScale, AdaScaleWrapper
label_desc: Use image node_id 4 for calling a global function with example usage: test_sequential(oss) without return types