code: def _test_grad_accum_func(rank, world_size, tempfile_name):
      _dist_init(
          rank, world_size, tempfile_name, backend="gloo"
      )  # Covers gloo
  
      model = Linear(4, 2, bias=False)
      model.to("cuda")
      model = DDP(model, device_ids=[rank])
      optim = AdaScale(
          SGD(model.parameters(), lr=0.1), num_gradients_to_accumulate=2
      )
      with model.no_sync():
          # iter 1, input vectors are pointing dim0 and dim1
          in_data = Tensor([0.0] * 4)
          in_data[rank] = 1.0
          in_data = in_data.cuda()
          out = model(in_data)
          out.sum().backward()
      # iter 2, input vectors are pointing dim2 and dim3
      in_data = Tensor([0.0] * 4)
      in_data[rank + 2] = 1.0
      in_data = in_data.cuda()
      out = model(in_data)
      out.sum().backward()
      # since all inputs are orthogonal, the gain should be exactly 4.0.
      assert np.allclose(optim.gain(), 4.0), optim.gain()
      optim.step()
      optim.zero_grad()
  
      dist.destroy_process_group()
  
imports: import tempfile
  import numpy
  import pytest
  import torch
  from torch import Tensor
  import torch.distributed
  import torch.multiprocessing
  from torch.nn import Linear
  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.internal import torch_version
  from fairscale.nn.data_parallel import FullyShardedDataParallel
  from fairscale.nn.data_parallel import ShardedDataParallel
  from fairscale.optim import OSS, AdaScale
label_desc: Use image node_id 5 for calling a global function with example usage: _test_grad_accum_func(rank, world_size, tempfile_name) without return types