code: def test_scale_not_equal_default(test_case):
      """Test gain value when scale doesn't equal world size * grad_accum"""
      scale = test_case["scale"]
      exp_gain = test_case["exp_gain"]
      model = Linear(4, 2, bias=False)
      optim = AdaScale(
          SGD(model.parameters(), lr=0.1),
          num_gradients_to_accumulate=4,
          scale=scale,
      )
  
      data = [
          [1.0, 0.0, 0.0, 0.0],
          [0.0, 1.0, 0.0, 0.0],
          [0.0, 0.0, 1.0, 0.0],
          [0.0, 0.0, 0.0, 1.0],
      ]
      for i in range(4):
          out = model(Tensor(data[i]))
          out.sum().backward()
      # Since the inputs are perfect orthogonal, the gain should be at the scale.
      assert np.allclose(optim.gain(), exp_gain), optim.gain()
  
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 10 for calling a global function with example usage: test_scale_not_equal_default(test_case) without return types