code: def test_add_param_group(debias_ewma, is_scaled_loss):
      """Test AdaScale supports add_param_group() API for both scaled and unscaled loss."""
      num_grads_to_accum = 2
      model1 = Linear(2, 2, bias=True)
      with torch.no_grad():
          # make weights and bias deterministic, which is needed for
          # multi-layer models. For them, adascale gain is affected by
          # parameters from other layers.
          model1.weight.copy_(
              Tensor([1.0, 2.0, 3.0, 4.0]).reshape(2, 2)
          )
          model1.bias.fill_(0.1)
      optim = AdaScale(
          SGD(model1.parameters(), lr=0.1),
          num_gradients_to_accumulate=2,
          is_scaled_loss=is_scaled_loss,
          debias_ewma=debias_ewma,
      )
      assert len(optim._hook_handles) == 2, len(optim._hook_handles)
  
      model2 = Linear(2, 3, bias=True)
      with torch.no_grad():
          # make weights and bias deterministic
          model2.weight.copy_(
              Tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).reshape(3, 2)
          )
          model2.bias.fill_(0.2)
      optim.add_param_group({"params": model2.parameters()})
      assert len(optim._hook_handles) == 4, len(optim._hook_handles)
  
      # make sure we can run the model.
      model = Sequential(model1, model2).cuda()
      in_data_0 = Tensor([1.0, 2.0]).cuda()
      loss = model(in_data_0).sum()
      if is_scaled_loss:
          loss = loss / num_grads_to_accum
      loss.backward()
  
      in_data_1 = Tensor([3.0, 4.0]).cuda()
      loss = model(in_data_1).sum()
      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()
  
      # make sure the gains are right and we can step.
      # since this is the first step, debias_ewma doesn't affect the value.
      assert np.allclose(optim.gain(), 1.1440223454935758), optim.gain()
      assert np.allclose(optim.gain(0), 1.1428571428571428), optim.gain(
          0
      )
      assert np.allclose(optim.gain(1), 1.1471258476157762), optim.gain(
          1
      )
      optim.step()
      optim.zero_grad()
  
      # make sure we can add a PG again after stepping.
      model3 = Linear(3, 4, bias=True)
      with torch.no_grad():
          # make weights and bias deterministic
          model3.weight.copy_(
              Tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0] * 2).reshape(4, 3)
          )
          model3.bias.fill_(0.2)
      optim.add_param_group({"params": model3.parameters()})
      assert len(optim._hook_handles) == 6, len(optim._hook_handles)
  
      # make sure we can run the model.
      model = Sequential(model1, model2, model3).cuda()
      in_data_0 = Tensor([1.0, 2.0]).cuda()
      loss = model(in_data_0).sum()
      if is_scaled_loss:
          loss = loss / num_grads_to_accum
      loss.backward()
  
      in_data_1 = Tensor([3.0, 4.0]).cuda()
      loss = model(in_data_1).sum()
      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()
  
      # make sure gains are right and we can step.
      # the last PG's gain is not affected by debias_ewma since it is the first step for that PG.
      assert np.allclose(
          optim.gain(),
          1.1382937715383077 if debias_ewma else 1.1391959826562015,
      ), optim.gain()
      assert np.allclose(
          optim.gain(0),
          1.142857206008338 if debias_ewma else 1.142857206006931,
      ), optim.gain(0)
      assert np.allclose(
          optim.gain(1),
          1.1116875516387468 if debias_ewma else 1.1116906378271827,
      ), optim.gain(1)
      assert np.allclose(optim.gain(2), 1.0749164095196344), optim.gain(
          2
      )
      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 6 for calling a global function with example usage: test_add_param_group(debias_ewma, is_scaled_loss) without return types