# Buggy code
def optimizer_closure(self, split_batch, batch_idx, opt_idx, optimizer, hiddens):
    """
    wrap the forward step in a closure so second order methods work
    """
    # ---------------------------
    # FORWARD
    # ---------------------------
    with self.profiler.profile("model_forward"):
        if self.use_amp and NATIVE_AMP_AVALAIBLE:
            with torch.cuda.amp.autocast():
                training_step_output = self.training_forward(
                    split_batch, batch_idx, opt_idx, hiddens
                )
        else:
            training_step_output = self.training_forward(
                split_batch, batch_idx, opt_idx, hiddens
            )

        # ----------------------------
        # PROCESS THE RESULT
        # ----------------------------
        # format and reduce outputs accordingly
        training_step_output = self.process_output(training_step_output, train=True)

        # TODO: temporary part of structured results PR
        training_step_output = AttributeDict(
            batch_loss=training_step_output[0],
            pbar_on_batch_end=training_step_output[1],
            log_metrics=training_step_output[2],
            callback_metrics=training_step_output[3],
            hiddens=training_step_output[4],
        )

        # if the user decides to finally reduce things in epoch_end, save raw output without graphs
        training_step_output_for_epoch_end = recursive_detach(training_step_output)

    # accumulate loss
    # (if accumulate_grad_batches = 1 no effect)
    closure_loss = training_step_output.batch_loss / self.accumulate_grad_batches

    # backward pass
    model_ref = self.get_model()
    with self.profiler.profile("model_backward"):
        # scale loss for 16 bit
        if self.precision == 16 and not self.on_tpu:
            closure_loss = model_ref.amp_scale_loss(closure_loss, optimizer, opt_idx)

        # do backward pass
        model_ref.backward(self, closure_loss, optimizer, opt_idx)

        # once backward has been applied, release graph
        closure_loss = closure_loss.detach()
        training_step_output.batch_loss = training_step_output.batch_loss.detach()

    if self.use_horovod:
        # Synchronize Horovod to ensure gradient manipulations (e.g., loss scaling) are valid
        optimizer.synchronize()

    # insert after step hook
    if self.is_function_implemented("on_after_backward"):
        model_ref = self.get_model()
        with self.profiler.profile("on_after_backward"):
            model_ref.on_after_backward()

    result = AttributeDict(
        loss=closure_loss,
        training_step_output=training_step_output,
        training_step_output_for_epoch_end=training_step_output_for_epoch_end,
        hiddens=training_step_output.hiddens,
    )
    return result


# Traceback (KeyError)
Traceback (most recent call last):
File "main.py", line 140, in <module>
main(hparams)
File "main.py", line 72, in main
trainer.fit(model)
File "/mnt/lustre/maxiao1/anaconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 881, in fit
self.ddp_train(task, model)
File "/mnt/lustre/maxiao1/anaconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/distrib_data_parallel.py", line 539, in ddp_train
self.run_pretrain_routine(model)
File "/mnt/lustre/maxiao1/anaconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1091, in run_pretrain_routine
self.train()
File "/mnt/lustre/maxiao1/anaconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 376, in train
self.run_training_epoch()
File "/mnt/lustre/maxiao1/anaconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 510, in run_training_epoch
self.run_training_epoch_end(epoch_output)
File "/mnt/lustre/maxiao1/anaconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/training_loop.py", line 535, in run_training_epoch_end
epoch_output = model.training_epoch_end(epoch_output)
File "/mnt/lustre/maxiao1/PVM/models/baseline.py", line 335, in training_epoch_end
avg_loss = torch.stack([x['loss'] for x in outputs]).mean()
File "/mnt/lustre/maxiao1/PVM/models/baseline.py", line 335, in <listcomp>
avg_loss = torch.stack([x['loss'] for x in outputs]).mean()
KeyError: 'loss'

# Fixed code
def optimizer_closure(self, split_batch, batch_idx, opt_idx, optimizer, hiddens):
    """
    wrap the forward step in a closure so second order methods work
    """
    # ---------------------------
    # FORWARD
    # ---------------------------
    with self.profiler.profile("model_forward"):
        if self.use_amp and NATIVE_AMP_AVALAIBLE:
            with torch.cuda.amp.autocast():
                training_step_output = self.training_forward(
                    split_batch, batch_idx, opt_idx, hiddens
                )
        else:
            training_step_output = self.training_forward(
                split_batch, batch_idx, opt_idx, hiddens
            )

        # ----------------------------
        # PROCESS THE RESULT
        # ----------------------------
        # format and reduce outputs accordingly
        training_step_output_for_epoch_end = training_step_output
        training_step_output = self.process_output(training_step_output, train=True)

        # TODO: temporary part of structured results PR
        training_step_output = AttributeDict(
            batch_loss=training_step_output[0],
            pbar_on_batch_end=training_step_output[1],
            log_metrics=training_step_output[2],
            callback_metrics=training_step_output[3],
            hiddens=training_step_output[4],
        )

        # if the user decides to finally reduce things in epoch_end, save raw output without graphs
        training_step_output_for_epoch_end = recursive_detach(
            training_step_output_for_epoch_end
        )

    # accumulate loss
    # (if accumulate_grad_batches = 1 no effect)
    closure_loss = training_step_output.batch_loss / self.accumulate_grad_batches

    # backward pass
    model_ref = self.get_model()
    with self.profiler.profile("model_backward"):
        # scale loss for 16 bit
        if self.precision == 16 and not self.on_tpu:
            closure_loss = model_ref.amp_scale_loss(closure_loss, optimizer, opt_idx)

        # do backward pass
        model_ref.backward(self, closure_loss, optimizer, opt_idx)

        # once backward has been applied, release graph
        closure_loss = closure_loss.detach()
        training_step_output.batch_loss = training_step_output.batch_loss.detach()

    if self.use_horovod:
        # Synchronize Horovod to ensure gradient manipulations (e.g., loss scaling) are valid
        optimizer.synchronize()

    # insert after step hook
    if self.is_function_implemented("on_after_backward"):
        model_ref = self.get_model()
        with self.profiler.profile("on_after_backward"):
            model_ref.on_after_backward()

    result = AttributeDict(
        loss=closure_loss,
        training_step_output=training_step_output,
        training_step_output_for_epoch_end=training_step_output_for_epoch_end,
        hiddens=training_step_output.hiddens,
    )
    return result