make example easily copy and runnable

Author: Alan Chu

docs/source-pytorch/common/tbptt.rst

@@ -10,6 +10,9 @@ hidden states should be kept in-between each time-dimension split.
 
 
 .. code-block:: python
+    import torch.optim as optim
+    import pytorch_lightning as pl
+    from pytorch_lightning import LightningModule
 
     class LitModel(LightningModule):
 
@@ -20,25 +23,33 @@ hidden states should be kept in-between each time-dimension split.
             self.automatic_optimization = False
 
             self.truncated_bptt_steps = 10
-            self.my_rnn = ...
+            self.my_rnn = ParityModuleRNN() # Define RNN model using ParityModuleRNN
 
         # 2. Remove the `hiddens` argument
         def training_step(self, batch, batch_idx):
 
             # 3. Split the batch in chunks along the time dimension
             split_batches = split_batch(batch, self.truncated_bptt_steps)
 
-            hiddens = ...  # 3. Choose the initial hidden state
+            hiddens = ... # Choose the initial hidden state
             for split_batch in range(split_batches):
                 # 4. Perform the optimization in a loop
                 loss, hiddens = self.my_rnn(split_batch, hiddens)
                 self.backward(loss)
-                optimizer.step()
-                optimizer.zero_grad()
+                self.optimizer.step()
+                self.optimizer.zero_grad()
 
                 # 5. "Truncate"
                 hiddens = hiddens.detach()
 
             # 6. Remove the return of `hiddens`
             # Returning loss in manual optimization is not needed
             return None
+
+        def configure_optimizers(self):
+            return optim.Adam(self.my_rnn.parameters(), lr=0.001)
+
+    if __name__ == "__main__":
+        model = LitModel()
+        trainer = pl.Trainer(max_epochs=5)
+        trainer.fit(model, train_dataloader) # Define your own dataloader