Repository: Lightning-AI/pytorch-lightning

Issue:
RuntimeError: OrderedDict mutated during iteration
## 🐛 Bug

I was getting `RuntimeError: OrderedDict mutated during iteration`.
It seems like using the same `LightningModule` object with `ModelSummary` and `Trainer` causes this error.

### To Reproduce
```python
from pytorch_lightning.core.memory import ModelSummary
model = CifarNet() # any pl module would work here
ModelSummary(model,mode='full')
trainer = Trainer(fast_dev_run=True,gpus=1)
trainer.fit(model)
```
Steps to reproduce the behavior:

1. View model summary using `ModelSummary` class
2. Call `trainer.fit` with same object.

### Stacktrace

```python
RuntimeError                              Traceback (most recent call last)
<ipython-input-20-8badc092c0ba> in <module>()
      1 # Checking for errors
      2 trainer = Trainer(fast_dev_run=True,gpus=1)
----> 3 trainer.fit(model)
11 frames
/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/trainer.py in fit(self, model, train_dataloader, val_dataloaders)
    916 
    917         elif self.single_gpu:
--> 918             self.single_gpu_train(model)
    919 
    920         elif self.use_tpu:  # pragma: no-cover
/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/distrib_parts.py in single_gpu_train(self, model)
    174             self.reinit_scheduler_properties(self.optimizers, self.lr_schedulers)
    175 
--> 176         self.run_pretrain_routine(model)
    177 
    178     def tpu_train(self, tpu_core_idx, model):
/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/trainer.py in run_pretrain_routine(self, model)
   1091 
   1092         # CORE TRAINING LOOP
-> 1093         self.train()
   1094 
   1095     def test(
/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py in train(self)
    373                 # RUN TNG EPOCH
    374                 # -----------------
--> 375                 self.run_training_epoch()
    376 
    377                 if self.max_steps and self.max_steps == self.global_step:
/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py in run_training_epoch(self)
    456             # RUN TRAIN STEP
    457             # ---------------
--> 458             _outputs = self.run_training_batch(batch, batch_idx)
    459             batch_result, grad_norm_dic, batch_step_metrics, batch_output = _outputs
    460 
/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py in run_training_batch(self, batch, batch_idx)
    632 
    633                 # calculate loss
--> 634                 loss, batch_output = optimizer_closure()
    635 
    636                 # check if loss or model weights are nan
/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py in optimizer_closure()
    596                                                                     opt_idx, self.hiddens)
    597                         else:
--> 598                             output_dict = self.training_forward(split_batch, batch_idx, opt_idx, self.hiddens)
    599 
    600                         # format and reduce outputs accordingly
/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/training_loop.py in training_forward(self, batch, batch_idx, opt_idx, hiddens)
    771             batch = self.transfer_batch_to_gpu(batch, gpu_id)
    772             args[0] = batch
--> 773             output = self.model.training_step(*args)
    774 
    775         # TPU support
<ipython-input-11-2482ebcf9d12> in training_step(self, batch, batch_idx)
     55   def training_step(self,batch,batch_idx):
     56     x, y = batch
---> 57     y_hat = self(x)
     58 
     59     return {'loss': F.cross_entropy(y_hat, y)}
/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    548             result = self._slow_forward(*input, **kwargs)
    549         else:
--> 550             result = self.forward(*input, **kwargs)
    551         for hook in self._forward_hooks.values():
    552             hook_result = hook(self, input, result)
<ipython-input-11-2482ebcf9d12> in forward(self, x)
     11 
     12   def forward(self,x):
---> 13     return self.model(x)
     14 
     15   def prepare_data(self):
/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    549         else:
    550             result = self.forward(*input, **kwargs)
--> 551         for hook in self._forward_hooks.values():
    552             hook_result = hook(self, input, result)
    553             if hook_result is not None:
RuntimeError: OrderedDict mutated during iteration
```

### Expected behavior
We should be able to use same object with both the classes

### Environment
```
* CUDA:
	- GPU:
		- Tesla T4
	- available:         True
	- version:           10.1
* Packages:
	- numpy:             1.18.5
	- pyTorch_debug:     False
	- pyTorch_version:   1.5.0+cu101
	- pytorch-lightning: 0.8.1
	- tensorboard:       2.2.2
	- tqdm:              4.41.1
* System:
	- OS:                Linux
	- architecture:
		- 64bit
		- 
	- processor:         x86_64
	- python:            3.6.9
	- version:           #1 SMP Wed Feb 19 05:26:34 PST 2020
```


Discussion:
Hi! thanks for your contribution!, great first issue!
> model = CifarNet() # any pl module would work here

Could you paste the minimal code for CifarNet? I cannot reproduce it with PL examples, sorry.
Okay ! I'm not sure which part is pertaining to this issue, so here is the link to my [colab notebook](https://colab.research.google.com/drive/13ER3opHF3IacEfAyEWojuplBUHn2MBKU?usp=sharing)

Resolving patch:
diff --git a/pytorch_lightning/core/memory.py b/pytorch_lightning/core/memory.py
--- a/pytorch_lightning/core/memory.py
+++ b/pytorch_lightning/core/memory.py
@@ -7,6 +7,7 @@
 import numpy as np
 import torch
 import torch.nn as nn
+from torch.utils.hooks import RemovableHandle
 
 import pytorch_lightning as pl
 from pytorch_lightning.utilities.apply_func import apply_to_collection
@@ -54,11 +55,17 @@ def __init__(self, module: nn.Module):
         self._in_size = None
         self._out_size = None
 
-    def _register_hook(self):
+    def __del__(self):
+        self.detach_hook()
+
+    def _register_hook(self) -> RemovableHandle:
         """
-        Registers a hook on the module that computes the input- and output size(s)
-        on the first forward pass. The hook will remove itself from the module, meaning that
+        Registers a hook on the module that computes the input- and output size(s) on the first forward pass.
+        If the hook is called, it will remove itself from the from the module, meaning that
         recursive models will only record their input- and output shapes once.
+
+        Return:
+            A handle for the installed hook.
         """
 
         def hook(module, inp, out):
@@ -66,16 +73,24 @@ def hook(module, inp, out):
                 inp = inp[0]
             self._in_size = parse_batch_shape(inp)
             self._out_size = parse_batch_shape(out)
-            self._hook_handle.remove()  # hook detaches itself from module
+            self._hook_handle.remove()
 
         return self._module.register_forward_hook(hook)
 
+    def detach_hook(self):
+        """
+        Removes the forward hook if it was not already removed in the forward pass.
+        Will be called after the summary is created.
+        """
+        if self._hook_handle is not None:
+            self._hook_handle.remove()
+
     @property
-    def in_size(self):
+    def in_size(self) -> Union[str, List]:
         return self._in_size or UNKNOWN_SIZE
 
     @property
-    def out_size(self):
+    def out_size(self) -> Union[str, List]:
         return self._out_size or UNKNOWN_SIZE
 
     @property
@@ -180,6 +195,8 @@ def summarize(self) -> Dict[str, LayerSummary]:
         summary = OrderedDict((name, LayerSummary(module)) for name, module in self.named_modules)
         if self._model.example_input_array is not None:
             self._forward_example_input()
+        for layer in summary.values():
+            layer.detach_hook()
         return summary
 
     def _forward_example_input(self) -> None:


Test patch:
diff --git a/tests/core/test_memory.py b/tests/core/test_memory.py
--- a/tests/core/test_memory.py
+++ b/tests/core/test_memory.py
@@ -89,6 +89,20 @@ def test_linear_model_summary_shapes(device, dtype, mode):
     assert model.device == device
 
 
+@pytest.mark.parametrize(['mode'], [
+    pytest.param(ModelSummary.MODE_FULL),
+    pytest.param(ModelSummary.MODE_TOP),
+])
+def test_hooks_removed_after_summarize(mode):
+    """ Test that all hooks were properly removed after summary, even ones that were not run. """
+    model = UnorderedModel()
+    summary = ModelSummary(model, mode=mode)
+    # hooks should be removed
+    for _, layer in summary.summarize().items():
+        handle = layer._hook_handle
+        assert handle.id not in handle.hooks_dict_ref()
+
+
 @pytest.mark.parametrize(['mode'], [
     pytest.param(ModelSummary.MODE_FULL),
     pytest.param(ModelSummary.MODE_TOP),