Update docstrings for backward methods

src/pytorch_lightning/plugins/precision/apex_amp.py

@@ -69,6 +69,7 @@ def backward(
             model: the model to be optimized
             closure_loss: the loss value obtained from the closure
             optimizer: current optimizer being used. ``None`` if using manual optimization
+            optimizer_idx: the index of the current optimizer. ``None`` if using manual optimization
         """
         opt = optimizer or model.trainer.optimizers
         with amp.scale_loss(closure_loss, opt) as closure_loss:

src/pytorch_lightning/plugins/precision/deepspeed.py

@@ -72,6 +72,16 @@ def backward(
         *args: Any,
         **kwargs: Any,
     ) -> None:
+        """Performs back-propagatin using DeepSpeed's engine.
+
+        Args:
+            model: the model to be optimized
+            closure_loss: the loss tensor
+            optimizer: ignored for DeepSpeed
+            optimizer_idx: ignored for DeepSpeed
+            *args: additional positional arguments for the :meth:`deepspeed.DeepSpeedEngine.backward` call
+            *kwargs: additional keyword arguments for the :meth:`deepspeed.DeepSpeedEngine.backward` call
+        """
         if is_overridden("backward", model):
             warning_cache.warn(
                 "You have overridden the `LightningModule.backward` hook but it will be ignored since DeepSpeed handles"

src/pytorch_lightning/plugins/precision/precision_plugin.py

@@ -74,6 +74,9 @@ def backward(
             model: the model to be optimized
             closure_loss: the loss value obtained from the closure
             optimizer: current optimizer being used. ``None`` if using manual optimization
+            *args: Positional arguments intended for the actual function that performs the backward, like
+                :meth:`~torch.Tensor.backward`.
+            *kwargs: Keyword arguments for the same purpose as `*args`.
         """
         # do backward pass
         if model is not None and isinstance(model, pl.LightningModule):

src/pytorch_lightning/strategies/strategy.py

@@ -183,6 +183,11 @@ def backward(
 
         Args:
             closure_loss: a tensor holding the loss value to backpropagate
+            optimizer: An optional optimizer that gets passed down to the precision plugin's backward
+            optimizer_idx: An optional optimizer index that gets passed down to the precision plugin's backward
+            *args: Positional arguments that get passed down to the precision plugin's backward, intended as arguments
+                for the actual function that performs the backward, like :meth:`~torch.Tensor.backward`.
+            *kwargs: Keyword arguments for the same purpose as `*args`.
         """
         self.pre_backward(closure_loss)
         assert self.lightning_module is not None