docstr: Check that a normal PyTorch ``lr_scheduler`` is usable with
  ZeroRedundancyOptimizer.
docstr_tok: ["Check","that","a","normal","PyTorch","``","lr_scheduler","``","is","usable","with","ZeroRedundancyOptimizer","."]
code: def test_lr_scheduler(self):
      """Check that a normal PyTorch ``lr_scheduler`` is usable with
      ZeroRedundancyOptimizer."""
      self.dist_init(self.rank)
      x = torch.tensor([1.0], device=self.device, requires_grad=True)
      x2 = torch.tensor([1.0], device=self.device, requires_grad=True)
      o = ZeroRedundancyOptimizer([x], optimizer_class=SGD, lr=0.01)
      o2 = torch.optim.SGD([x2], lr=0.01)
      s = torch.optim.lr_scheduler.StepLR(o, 1)
      s2 = torch.optim.lr_scheduler.StepLR(o2, 1)
      for _ in range(5):
          x.backward()
          o.zero_grad()
          o.step()
          s.step()
          x2.backward()
          o2.zero_grad()
          o2.step()
          s2.step()
          self.assertEqual(x, x2)
  
imports: import copy
  import os
  import sys
  import unittest
  from contextlib import nullcontext
  from typing import Any, cast, List
  import numpy
  import torch
  import torch.distributed
  from torch.distributed.algorithms.ddp_comm_hooks.ddp_zero_hook import hook_with_zero_step, hook_with_zero_step_interleaved
  from torch.distributed.algorithms.ddp_comm_hooks.default_hooks import allreduce_hook
  from torch.distributed.algorithms.join import Join, Joinable, JoinHook
  from torch.distributed.optim import ZeroRedundancyOptimizer
  from torch.distributed.optim.zero_redundancy_optimizer import _broadcast_object
  from torch.nn.parallel import DistributedDataParallel
  from torch.optim import AdamW, SGD
  from torch.testing._internal import common_distributed
  from torch.testing._internal.common_utils import instantiate_parametrized_tests, IS_WINDOWS, parametrize, run_tests, TEST_WITH_ASAN, TEST_WITH_DEV_DBG_ASAN
label_desc: Use image node_id 7 for calling the TestZeroRedundancyOptimizerDistributed obj's underlying member method code with example usage: obj.test_lr_scheduler() without return types