add python scalars to docstring

Lightning-AI · Jul 28, 2020 · 82d853b · 82d853b
1 parent 1305a34
commit 82d853b
Showing 1 changed file with 9 additions and 9 deletions.
diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
@@ -167,7 +167,7 @@ def forward(self, batch):
 
         """
 
-    def training_step(self, *args, **kwargs) -> Union[int, Dict[str, Union[Tensor, Dict[str, Tensor]]]]:
+    def training_step(self, *args, **kwargs) -> Union[int, Dict[str, Union[Tensor, Dict[str, Union[float, Tensor]]]]]:
         r"""
         Here you compute and return the training loss and some additional metrics for e.g.
         the progress bar or logger.
@@ -185,8 +185,8 @@ def training_step(self, *args, **kwargs) -> Union[int, Dict[str, Union[Tensor, D
             When implementing :meth:`training_step`, return whatever you need in that step:
 
             - loss -> tensor scalar **REQUIRED**
-            - progress_bar -> Dict for progress bar display. Must have either tensors or scalars
-            - log -> Dict of metrics to add to logger. Must have either tensors or scalars (no images, etc)
+            - progress_bar -> Dict for progress bar display. Must have either scalar tensors or Python scalars
+            - log -> Dict of metrics to add to logger. Must have either scalar tensors or Python scalars (no images, etc)
 
         In this step you'd normally do the forward pass and calculate the loss for a batch.
         You can also do fancier things like multiple forward passes or something model specific.
@@ -357,8 +357,8 @@ def training_step_end(self, *args, **kwargs) -> Dict[str, Union[Tensor, Dict[str
             Dict with loss key and optional log or progress bar keys.
 
             - loss -> tensor scalar **REQUIRED**
-            - progress_bar -> Dict for progress bar display. Must have only tensors
-            - log -> Dict of metrics to add to logger. Must have only tensors (no images, etc)
+            - progress_bar -> Dict for progress bar display. Must have either scalar tensors or Python scalars
+            - log -> Dict of metrics to add to logger. Must have either scalar tensors or Python scalars (no images, etc)
 
         Examples:
             .. code-block:: python
@@ -574,8 +574,8 @@ def validation_epoch_end(
             Dict or OrderedDict.
             May have the following optional keys:
 
-            - progress_bar (dict for progress bar display; either tensors or scalars)
-            - log (dict of metrics to add to logger; either tensors or scalars).
+            - progress_bar (dict for progress bar display; either scalar tensors or Python scalars)
+            - log (dict of metrics to add to logger; either scalar tensors or Python scalars).
 
         Note:
             If you didn't define a :meth:`validation_step`, this won't be called.
@@ -799,8 +799,8 @@ def test_epoch_end(
         Return:
             Dict or OrderedDict: Dict has the following optional keys:
 
-            - progress_bar -> Dict for progress bar display. Must have only tensors.
-            - log -> Dict of metrics to add to logger. Must have only tensors (no images, etc).
+            - progress_bar -> Dict for progress bar display. Must have either scalar tensors or Python scalars.
+            - log -> Dict of metrics to add to logger. Must have either scalar tensors or Python scalars (no images, etc).
 
         Note:
             If you didn't define a :meth:`test_step`, this won't be called.