docstr: Check that the ``step(**kwargs)`` interface is properly exposed.
docstr_tok: ["Check","that","the","``","step","(","*","*","kwargs",")","``","interface","is","properly","exposed","."]
code: def test_step_with_kwargs(self):
      """Check that the ``step(**kwargs)`` interface is properly exposed."""
      self.dist_init(self.rank)
      LR = 0.1
  
      class SGDWithStepKWArg(torch.optim.SGD):
          def step(self, closure=None, kwarg=None):
              super().step()
              kwarg.append(5)
  
      kwarg: List[Any] = []
      x = torch.tensor([1.0], device=self.device, requires_grad=True)
      o = ZeroRedundancyOptimizer(
          [x],
          optimizer_class=SGDWithStepKWArg,
          lr=LR,
      )
      x.backward()
      o.step(0, kwarg=kwarg)
      self.assertEqual(kwarg, [5])
      self.assertEqual(x, torch.tensor([0.9], device=self.device))
  
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 3 for calling the TestZeroRedundancyOptimizerSingleRank obj's underlying member method code with example usage: obj.test_step_with_kwargs() without return types