Commit f5cd5ae9 authored by Michael Carilli's avatar Michael Carilli
Browse files

Minor docstring updates

parent e7f19560
...@@ -37,7 +37,7 @@ def scale_loss(loss, ...@@ -37,7 +37,7 @@ def scale_loss(loss,
unscaled. The direct ``.grad`` attributes of any FP16 unscaled. The direct ``.grad`` attributes of any FP16
model params will remain scaled after context manager exit. model params will remain scaled after context manager exit.
This subtlety affects gradient clipping. See "Gradient clipping" under This subtlety affects gradient clipping. See "Gradient clipping" under
`Advanced use cases`_ for best practices. `Advanced Amp Usage`_ for best practices.
Args: Args:
loss(Tensor): Typically a scalar Tensor. The ``scaled_loss`` that the context loss(Tensor): Typically a scalar Tensor. The ``scaled_loss`` that the context
...@@ -48,16 +48,17 @@ def scale_loss(loss, ...@@ -48,16 +48,17 @@ def scale_loss(loss,
model(torch.nn.Module, optional, default=None): Currently unused, reserved to enable future model(torch.nn.Module, optional, default=None): Currently unused, reserved to enable future
optimizations. optimizations.
delay_unscale(bool, default=False): Don't unscale the gradients or perform model->master delay_unscale(bool, default=False): Don't unscale the gradients or perform model->master
gradient copies on context manager exit. "Advanced use cases" illustrates gradient copies on context manager exit. `Advanced Amp Usage`_ illustrates
situations where this is necessary. situations where this is necessary.
.. warning::If ``True``, ``optimizer.step()`` cannot be .. warning::
called yet after context manager exit, and must wait for another, later backward context If ``delay_unscale`` is ``True`` for a given backward pass, ``optimizer.step()`` cannot be
manager invocation with ``delay_unscale`` left to False. called yet after context manager exit, and must wait for another, later backward context
See `Advanced use cases`_ for examples. manager invocation with ``delay_unscale`` left to False.
See `Advanced Amp Usage`_ for examples.
.. _`Advanced use cases`: .. _`Advanced Amp Usage`:
https://nvidia.github.io/apex/amp.html#advanced-use-cases https://nvidia.github.io/apex/advanced.html
""" """
if not _amp_state.opt_properties.enabled: if not _amp_state.opt_properties.enabled:
yield loss yield loss
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment