parameters: (fwd=None, **kwargs)
return_statement: return decorate_fwd
docstring: Helper decorator for ``forward`` methods of custom autograd functions (subclasses of
      :class:`torch.autograd.Function`).  See the :ref:`example page<amp-custom-examples>` for more detail.
  
      Args:
          cast_inputs (:class:`torch.dtype` or None, optional, default=None):  If not ``None``,
              when ``forward`` runs in an autocast-enabled region, casts incoming
              floating-point CUDA Tensors to the target dtype (non-floating-point Tensors are not affected),
              then executes ``forward`` with autocast disabled.
              If ``None``, ``forward``'s internal ops execute with the current autocast state.
  
      .. note::
          If the decorated ``forward`` is called outside an autocast-enabled region,
          :func:`custom_fwd<custom_fwd>` is a no-op and ``cast_inputs`` has no effect.
docstring_summary: Helper decorator for ``forward`` methods of custom autograd functions (subclasses of
      :class:`torch.autograd.Function`).  See the :ref:`example page<amp-custom-examples>` for more detail.
docstring_tokens: Helper decorator for forward methods of custom autograd functions ( subclasses of : class : torch . autograd . Function ) . See the : ref : example page<amp - custom - examples > for more detail .
function: def custom_fwd(fwd=None, **kwargs):
      """
      Helper decorator for ``forward`` methods of custom autograd functions (subclasses of
      :class:`torch.autograd.Function`).  See the :ref:`example page<amp-custom-examples>` for more detail.
  
      Args:
          cast_inputs (:class:`torch.dtype` or None, optional, default=None):  If not ``None``,
              when ``forward`` runs in an autocast-enabled region, casts incoming
              floating-point CUDA Tensors to the target dtype (non-floating-point Tensors are not affected),
              then executes ``forward`` with autocast disabled.
              If ``None``, ``forward``'s internal ops execute with the current autocast state.
  
      .. note::
          If the decorated ``forward`` is called outside an autocast-enabled region,
          :func:`custom_fwd<custom_fwd>` is a no-op and ``cast_inputs`` has no effect.
      """
      if fwd is None:
          if len(kwargs) == 0:
              cast_inputs = None
          else:
              assert len(kwargs) == 1
              cast_inputs = kwargs["cast_inputs"]
          return functools.partial(custom_fwd, cast_inputs=cast_inputs)
  
      if len(kwargs) == 0:
          cast_inputs = None
      else:
          assert len(kwargs) == 1
          cast_inputs = kwargs["cast_inputs"]
  
      @functools.wraps(fwd)
      def decorate_fwd(*args, **kwargs):
          if cast_inputs is None:
              args[0]._fwd_used_autocast = torch.is_autocast_enabled()
              return fwd(*args, **kwargs)
          else:
              autocast_context = torch.is_autocast_enabled()
              args[0]._fwd_used_autocast = False
              if autocast_context:
                  with autocast(enabled=False):
                      return fwd(*_cast(args, cast_inputs), **_cast(kwargs, cast_inputs))
              else:
                  return fwd(*args, **kwargs)
      return decorate_fwd
function_tokens: def custom_fwd ( fwd = None , * * kwargs ) : if fwd is None : if len ( kwargs ) == 0 : cast_inputs = None else : assert len ( kwargs ) == 1 cast_inputs = kwargs [ "cast_inputs" ] return functools . partial ( custom_fwd , cast_inputs = cast_inputs ) if len ( kwargs ) == 0 : cast_inputs = None else : assert len ( kwargs ) == 1 cast_inputs = kwargs [ "cast_inputs" ] @ functools . wraps ( fwd ) def decorate_fwd ( * args , * * kwargs ) : if cast_inputs is None : args [ 0 ] . _fwd_used_autocast = torch . is_autocast_enabled ( ) return fwd ( * args , * * kwargs ) else : autocast_context = torch . is_autocast_enabled ( ) args [ 0 ] . _fwd_used_autocast = False if autocast_context : with autocast ( enabled = False ) : return fwd ( * _cast ( args , cast_inputs ) , * * _cast ( kwargs , cast_inputs ) ) else : return fwd ( * args , * * kwargs ) return decorate_fwd