convert doctest to code block

Lightning-AI · Jun 8, 2020 · 9b78794 · 9b78794
1 parent c72b709
commit 9b78794
Showing 1 changed file with 18 additions and 9 deletions.
diff --git a/pytorch_lightning/core/decorators.py b/pytorch_lightning/core/decorators.py
@@ -34,15 +34,24 @@ def auto_move_data(fn: Callable) -> Callable:
 
     Example:
 
-        >>> class LitModel(LightningModule):
-        ...     @auto_move_data
-        ...     def forward(self, x):
-        ...         return x
-        >>> LitModel.forward = auto_move_data(LitModel.forward)
-        >>> model = LitModel()
-        >>> model = model.to('cuda')
-        >>> model(torch.zeros(1, 3))
-        tensor([[0., 0., 0.]], device='cuda:0')
+        .. code-block:: python
+
+            # directly in the source code
+            class LitModel(LightningModule):
+
+                @auto_move_data
+                def forward(self, x):
+                    return x
+
+            # or outside
+            LitModel.forward = auto_move_data(LitModel.forward)
+
+            model = LitModel()
+            model = model.to('cuda')
+            model(torch.zeros(1, 3))
+
+            # input gets moved to device
+            # tensor([[0., 0., 0.]], device='cuda:0')
 
     """
     @wraps(fn)