# Buggy code
def _register_hook(self):
    """
    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
    recursive models will only record their input- and output shapes once.
    """

    def hook(module, inp, out):
        if len(inp) == 1:
            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

    return self._module.register_forward_hook(hook)


# 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 _register_hook(self) -> RemovableHandle:
    """
    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):
        if len(inp) == 1:
            inp = inp[0]
        self._in_size = parse_batch_shape(inp)
        self._out_size = parse_batch_shape(out)
        self._hook_handle.remove()

    return self._module.register_forward_hook(hook)