# Buggy code
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()
    return summary


# Traceback (RuntimeError)
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

# Fixed code
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