docstr: Check that ZeroRedundancyOptimizer properly exposes the ``step()``
  interface.
docstr_tok: ["Check","that","ZeroRedundancyOptimizer","properly","exposes","the","``","step","(",")","``","interface","."]
code: def test_step(self):
      """Check that ZeroRedundancyOptimizer properly exposes the ``step()``
      interface."""
      self.dist_init(self.rank, world_size=self.world_size)
      LR = 0.01
  
      with self.context:
          x = torch.tensor([float(self.rank + 1)], device=self.device)
          m = torch.nn.Linear(1, 1)
          m.weight.data = torch.tensor([[1.0]])
          m.bias.data = torch.tensor([2.0])
          m = m.to(self.device)
          m_zero = copy.deepcopy(m).to(self.device)
  
          o = SGD(m.parameters(), lr=LR)
          o_zero = ZeroRedundancyOptimizer(
              m_zero.parameters(),
              optimizer_class=SGD,
              lr=LR,
          )
  
          y = m(x)
          y.backward(x)
          y_zero = m_zero(x)
          y_zero.backward(x)
  
          for p in m.parameters():
              dist.all_reduce(p.grad.data, op=dist.ReduceOp.SUM)
              p.grad.data /= self.world_size
          o.step()
          for p in m_zero.parameters():
              dist.all_reduce(p.grad.data, op=dist.ReduceOp.SUM)
              p.grad.data /= self.world_size
          o_zero.step()
  
          self.assertEqual(m.weight, m_zero.weight)
          self.assertEqual(m.bias, m_zero.bias)
  
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 5 for calling the TestZeroRedundancyOptimizerDistributed obj's underlying member method code with example usage: obj.test_step() without return types