code: def test_set_num_gradients_to_accumulate(test_case, is_scaled_loss):
      """Test set_num_gradients_to_accumulate experimental feature."""
      num_grads_to_accum = test_case["num_grads_to_accum"]
      exp_gain = test_case["exp_gain"]
  
      model = Linear(2, 2, bias=False)
      optim = AdaScale(
          SGD(model.parameters(), lr=0.1),
          num_gradients_to_accumulate=2,
          is_scaled_loss=is_scaled_loss,
      )
      loss = model(Tensor([0.0, 1.0])).sum()
      if is_scaled_loss:
          loss = loss / 2
      loss.backward()
      loss = model(Tensor([1.0, 0.0])).sum()
      if is_scaled_loss:
          loss = loss / 2
      loss.backward()
  
      if not is_scaled_loss:
          optim.scale_grad_by_num_grads_to_accum()
  
      assert np.allclose(optim.gain(), 2.0)
      optim.step()
      optim.zero_grad()
  
      optim.set_scale(float(num_grads_to_accum))
      optim.set_num_gradients_to_accumulate(num_grads_to_accum)
      for _ in range(num_grads_to_accum):
          loss = model(Tensor([0.0, 1.0])).sum() / num_grads_to_accum
          if is_scaled_loss:
              loss = loss / num_grads_to_accum
          loss.backward()
  
      if not is_scaled_loss:
          optim.scale_grad_by_num_grads_to_accum()
  
      assert np.allclose(optim.gain(), exp_gain), optim.gain()
      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 7 for calling a global function with example usage: test_set_num_gradients_to_accumulate(test_case, is_scaled_loss) without return types