code: def test_lr_scheduler():
      """Test AdaScale working with torch.optim.lr_scheduler."""
      num_grads_to_accum = 3
      model = Linear(2, 2, bias=False)
      optim = AdaScale(
          SGD(model.parameters(), lr=0.1),
          num_gradients_to_accumulate=num_grads_to_accum,
      )
      # We use 1, not 0.1 here since scheduler.step() is called here first.
      scheduler = LambdaLR(
          optim, lr_lambda=lambda epoch: 1 / 10**epoch
      )
      for epoch in range(3):
          for data_idx in range(10):
              for accumulation in range(num_grads_to_accum):
                  in_data = torch.rand(2)
                  loss = model(in_data).sum()
                  loss.backward()
              assert optim.gain() <= 3, optim.gain()
              optim.step()
              optim.zero_grad()
              # asserting LR is right
              assert np.allclose(
                  optim.param_groups[0]["lr"], 0.1 / 10**epoch
              ), optim.param_groups[0]["lr"]
          scheduler.step()
          # asserting LR is right
          assert np.allclose(
              optim.param_groups[0]["lr"], 0.1 / 10 ** (epoch + 1)
          ), optim.param_groups[0]["lr"]
  
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 5 for calling a global function with example usage: test_lr_scheduler() without return types