parameters: (self, optimizer, *args, **kwargs)
return_statement: return retval
docstring: :meth:`step` carries out the following two operations:
  
          1.  Internally invokes ``unscale_(optimizer)`` (unless :meth:`unscale_` was explicitly called for ``optimizer``
              earlier in the iteration).  As part of the :meth:`unscale_`, gradients are checked for infs/NaNs.
          2.  If no inf/NaN gradients are found, invokes ``optimizer.step()`` using the unscaled
              gradients.  Otherwise, ``optimizer.step()`` is skipped to avoid corrupting the params.
  
          ``*args`` and ``**kwargs`` are forwarded to ``optimizer.step()``.
  
          Returns the return value of ``optimizer.step(*args, **kwargs)``.
  
          Args:
              optimizer (torch.optim.Optimizer):  Optimizer that applies the gradients.
              args:  Any arguments.
              kwargs:  Any keyword arguments.
  
          .. warning::
              Closure use is not currently supported.
docstring_summary: :meth:`step` carries out the following two operations:
docstring_tokens: : meth : step carries out the following two operations :
function: def step(self, optimizer, *args, **kwargs):
          """
          :meth:`step` carries out the following two operations:
  
          1.  Internally invokes ``unscale_(optimizer)`` (unless :meth:`unscale_` was explicitly called for ``optimizer``
              earlier in the iteration).  As part of the :meth:`unscale_`, gradients are checked for infs/NaNs.
          2.  If no inf/NaN gradients are found, invokes ``optimizer.step()`` using the unscaled
              gradients.  Otherwise, ``optimizer.step()`` is skipped to avoid corrupting the params.
  
          ``*args`` and ``**kwargs`` are forwarded to ``optimizer.step()``.
  
          Returns the return value of ``optimizer.step(*args, **kwargs)``.
  
          Args:
              optimizer (torch.optim.Optimizer):  Optimizer that applies the gradients.
              args:  Any arguments.
              kwargs:  Any keyword arguments.
  
          .. warning::
              Closure use is not currently supported.
          """
          if (not self._enabled):
              return optimizer.step(*args, **kwargs)
  
          if "closure" in kwargs:
              raise RuntimeError("Closure use is not currently supported if GradScaler is enabled.")
  
          self._check_scale_growth_tracker("step")
  
          optimizer_state = self._per_optimizer_states[id(optimizer)]
  
          if optimizer_state["stage"] is OptState.STEPPED:
              raise RuntimeError("step() has already been called since the last update().")
  
          retval = None
  
          if (hasattr(optimizer, "_step_supports_amp_scaling") and optimizer._step_supports_amp_scaling):
              # This optimizer has customized scale-handling logic, so we can call optimizer.step() directly.
              # The contract with custom optimizers is that their step() should accept an additional,
              # optional grad_scaler kwarg.  We append self to the kwargs so the custom optimizer has full information:
              # it can query its own state, invoke unscale_ on itself, etc
              retval = optimizer.step(*args, **dict(kwargs, grad_scaler=self))
              optimizer_state["stage"] = OptState.STEPPED
              return retval
  
          if optimizer_state["stage"] is OptState.READY:
              self.unscale_(optimizer)
  
          assert len(optimizer_state["found_inf_per_device"]) > 0, "No inf checks were recorded for this optimizer."
  
          retval = self._maybe_opt_step(optimizer, optimizer_state, *args, **kwargs)
  
          optimizer_state["stage"] = OptState.STEPPED
  
          return retval
function_tokens: def step ( self , optimizer , * args , * * kwargs ) : if ( not self . _enabled ) : return optimizer . step ( * args , * * kwargs ) if "closure" in kwargs : raise RuntimeError ( "Closure use is not currently supported if GradScaler is enabled." ) self . _check_scale_growth_tracker ( "step" ) optimizer_state = self . _per_optimizer_states [ id ( optimizer ) ] if optimizer_state [ "stage" ] is OptState . STEPPED : raise RuntimeError ( "step() has already been called since the last update()." ) retval = None if ( hasattr ( optimizer , "_step_supports_amp_scaling" ) and optimizer . _step_supports_amp_scaling ) : # This optimizer has customized scale-handling logic, so we can call optimizer.step() directly. # The contract with custom optimizers is that their step() should accept an additional, # optional grad_scaler kwarg.  We append self to the kwargs so the custom optimizer has full information: # it can query its own state, invoke unscale_ on itself, etc retval = optimizer . step ( * args , * * dict ( kwargs , grad_scaler = self ) ) optimizer_state [ "stage" ] = OptState . STEPPED return retval if optimizer_state [ "stage" ] is OptState . READY : self . unscale_ ( optimizer ) assert len ( optimizer_state [ "found_inf_per_device" ] ) > 0 , "No inf checks were recorded for this optimizer." retval = self . _maybe_opt_step ( optimizer , optimizer_state , * args , * * kwargs ) optimizer_state [ "stage" ] = OptState . STEPPED return retval