code: def test_grad_accum(test_case, cpu, is_scaled_loss):
      """Test the basic functionality on CPU/GPU with gradient accumulation without DDP"""
      make_cudnn_deterministic()
      model = Linear(2, 2, bias=True)
      if not cpu:
          if torch.cuda.device_count() < 1:
              pytest.skip("1 GPU is required")
          model = model.cuda()
      optim = AdaScale(
          SGD(model.parameters(), lr=0.1),
          num_gradients_to_accumulate=2,
          is_scaled_loss=is_scaled_loss,
      )
      expected_gain = test_case["expected_gain"]
      if "input" in test_case:
          data = [test_case["input"]] * 2
          gains = [expected_gain] * 2
      else:
          data = test_case["inputs"]
          gains = [None, expected_gain]
      for in_data, exp_gain in zip(
          data, gains
      ):  # test 2 iterations catch more corner cases.
          # grad pass 1
          in_data_0 = Tensor(in_data[0])
          if not cpu:
              in_data_0 = in_data_0.cuda()
          loss = model(in_data_0).sum()
          if is_scaled_loss:
              loss = loss / 2
          loss.backward()
          # grad pass 2
          in_data_1 = Tensor(in_data[1])
          if not cpu:
              in_data_1 = in_data_1.cuda()
          loss = model(in_data_1).sum()
          if is_scaled_loss:
              loss = loss / 2
          loss.backward()
  
          if not is_scaled_loss:
              optim.scale_grad_by_num_grads_to_accum()
  
          if exp_gain is not None:
              assert np.allclose(optim.gain(), exp_gain), optim.gain()
              w, b = model.parameters()
              assert np.allclose(
                  w.grad.cpu(), test_case["expected_grad"]
              ), w.grad
              assert np.allclose(
                  b.grad.cpu(), test_case["expected_bias_grad"]
              ), b.grad
          # stepping it. Note that if we did more than 2 passes as promised by the
          # num_gradients_to_accumulate argument above, AdaScale is not be able to
          # detect that mistake for now. The result will just be wrong in that case.
          optim.step()
          optim.zero_grad()
  
imports: import tempfile
  import numpy
  import pytest
  import torch
  from torch import Tensor
  from torch.nn import Linear, Sequential
  from torch.optim import SGD
  from torch.optim.lr_scheduler import LambdaLR
  from fairscale.fair_dev.testing.golden_testing_data import adascale_test_data
  from fairscale.fair_dev.testing.testing import make_cudnn_deterministic, skip_if_no_cuda
  from fairscale.fair_dev.testing.testing_memory import find_tensor_by_shape
  from fairscale.optim import AdaScale
label_desc: Use image node_id 3 for calling a global function with example usage: test_grad_accum(test_case, cpu, is_scaled_loss) without return types