# Buggy code
def process_output(self, output, train=False):
    """Reduces output according to the training mode.

    Separates loss from logging and progress bar metrics
    """
    # --------------------------
    # handle single scalar only
    # --------------------------
    # single scalar returned from a xx_step
    if isinstance(output, torch.Tensor):
        progress_bar_metrics = {}
        log_metrics = {}
        callback_metrics = {}
        hiddens = None
        return output, progress_bar_metrics, log_metrics, callback_metrics, hiddens

    # ---------------
    # EXTRACT CALLBACK KEYS
    # ---------------
    # all keys not progress_bar or log are candidates for callbacks
    callback_metrics = {}
    for k, v in output.items():
        if k not in ["progress_bar", "log", "hiddens"]:
            callback_metrics[k] = v

    if train and (self.use_dp or self.use_ddp2):
        num_gpus = self.num_gpus
        callback_metrics = self.reduce_distributed_output(callback_metrics, num_gpus)

    # ---------------
    # EXTRACT PROGRESS BAR KEYS
    # ---------------
    try:
        progress_output = output["progress_bar"]

        # reduce progress metrics for progress bar when using dp
        if train and (self.use_dp or self.use_ddp2):
            num_gpus = self.num_gpus
            progress_output = self.reduce_distributed_output(progress_output, num_gpus)

        progress_bar_metrics = progress_output
    except Exception:
        progress_bar_metrics = {}

    # ---------------
    # EXTRACT LOGGING KEYS
    # ---------------
    # extract metrics to log to experiment
    try:
        log_output = output["log"]

        # reduce progress metrics for progress bar when using dp
        if train and (self.use_dp or self.use_ddp2):
            num_gpus = self.num_gpus
            log_output = self.reduce_distributed_output(log_output, num_gpus)

        log_metrics = log_output
    except Exception:
        log_metrics = {}

    # ---------------
    # EXTRACT LOSS
    # ---------------
    # if output dict doesn't have the keyword loss
    # then assume the output=loss if scalar
    loss = None
    if train:
        try:
            loss = output["loss"]
        except Exception:
            if isinstance(output, torch.Tensor):
                loss = output
            else:
                raise RuntimeError(
                    "No `loss` value in the dictionary returned from `model.training_step()`."
                )

        # when using dp need to reduce the loss
        if self.use_dp or self.use_ddp2:
            loss = self.reduce_distributed_output(loss, self.num_gpus)

    # ---------------
    # EXTRACT HIDDEN
    # ---------------
    hiddens = output.get("hiddens")

    # use every metric passed in as a candidate for callback
    callback_metrics.update(progress_bar_metrics)
    callback_metrics.update(log_metrics)

    # detach all metrics for callbacks to prevent memory leaks
    # no .item() because it will slow things down
    callback_metrics = recursive_detach(callback_metrics)

    return loss, progress_bar_metrics, log_metrics, callback_metrics, hiddens


# Traceback (AttributeError)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-183-013dcacba267> in <module>()
10 # train
11 trainer = pl.Trainer(progress_bar_refresh_rate=0, max_epochs=5, weights_summary=None, gpus=1)
---> 12 trainer.fit(model, train_loader, val_loader)

7 frames
/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/states.py in wrapped_fn(self, *args, **kwargs)
46             if entering is not None:
47                 self.state = entering
---> 48             result = fn(self, *args, **kwargs)
49
50             # The INTERRUPTED state can be set inside the run function. To indicate that run was interrupted

/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/trainer.py in fit(self, model, train_dataloader, val_dataloaders, datamodule)
1071             self.accelerator_backend = GPUBackend(self)
1072             model = self.accelerator_backend.setup(model)
-> 1073             results = self.accelerator_backend.train(model)
1074
1075         elif self.use_tpu:

/usr/local/lib/python3.6/dist-packages/pytorch_lightning/accelerators/gpu_backend.py in train(self, model)
49
50     def train(self, model):
---> 51         results = self.trainer.run_pretrain_routine(model)
52         return results
53

/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/trainer.py in run_pretrain_routine(self, model)
1222
1223         # run a few val batches before training starts
-> 1224         self._run_sanity_check(ref_model, model)
1225
1226         # clear cache before training

/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/trainer.py in _run_sanity_check(self, ref_model, model)
1255             num_loaders = len(self.val_dataloaders)
1256             max_batches = [self.num_sanity_val_steps] * num_loaders
-> 1257             eval_results = self._evaluate(model, self.val_dataloaders, max_batches, False)
1258
1259             # allow no returns from eval

/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/evaluation_loop.py in _evaluate(self, model, dataloaders, max_batches, test_mode)
397
398         # log callback metrics
--> 399         self.__update_callback_metrics(eval_results, using_eval_result)
400
401         # Write predictions to disk if they're available.

/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/evaluation_loop.py in __update_callback_metrics(self, eval_results, using_eval_result)
429                         flat = {'val_loss': eval_result}
430                     else:
--> 431                         flat = flatten_dict(eval_result)
432                     self.callback_metrics.update(flat)
433             else:

/usr/local/lib/python3.6/dist-packages/pytorch_lightning/utilities/parsing.py in flatten_dict(source, result)
113         result = {}
114
--> 115     for k, v in source.items():
116         if isinstance(v, dict):
117             _ = flatten_dict(v, result)

AttributeError: 'NoneType' object has no attribute 'items'

# Fixed code
def process_output(self, output, train=False):
    """Reduces output according to the training mode.

    Separates loss from logging and progress bar metrics
    """
    # --------------------------
    # handle single scalar only
    # --------------------------
    # single scalar returned from a xx_step
    if isinstance(output, torch.Tensor):
        progress_bar_metrics = {}
        log_metrics = {}
        callback_metrics = {}
        hiddens = None
        return output, progress_bar_metrics, log_metrics, callback_metrics, hiddens

    # ---------------
    # EXTRACT CALLBACK KEYS
    # ---------------
    # all keys not progress_bar or log are candidates for callbacks
    callback_metrics = {}
    if output:
        for k, v in output.items():
            if k not in ["progress_bar", "log", "hiddens"]:
                callback_metrics[k] = v

    if train and (self.use_dp or self.use_ddp2):
        num_gpus = self.num_gpus
        callback_metrics = self.reduce_distributed_output(callback_metrics, num_gpus)

    # ---------------
    # EXTRACT PROGRESS BAR KEYS
    # ---------------
    try:
        progress_output = output["progress_bar"]

        # reduce progress metrics for progress bar when using dp
        if train and (self.use_dp or self.use_ddp2):
            num_gpus = self.num_gpus
            progress_output = self.reduce_distributed_output(progress_output, num_gpus)

        progress_bar_metrics = progress_output
    except Exception:
        progress_bar_metrics = {}

    # ---------------
    # EXTRACT LOGGING KEYS
    # ---------------
    # extract metrics to log to experiment
    try:
        log_output = output["log"]

        # reduce progress metrics for progress bar when using dp
        if train and (self.use_dp or self.use_ddp2):
            num_gpus = self.num_gpus
            log_output = self.reduce_distributed_output(log_output, num_gpus)

        log_metrics = log_output
    except Exception:
        log_metrics = {}

    # ---------------
    # EXTRACT LOSS
    # ---------------
    # if output dict doesn't have the keyword loss
    # then assume the output=loss if scalar
    loss = None
    if train:
        try:
            loss = output["loss"]
        except Exception:
            if isinstance(output, torch.Tensor):
                loss = output
            else:
                raise RuntimeError(
                    "No `loss` value in the dictionary returned from `model.training_step()`."
                )

        # when using dp need to reduce the loss
        if self.use_dp or self.use_ddp2:
            loss = self.reduce_distributed_output(loss, self.num_gpus)

    # ---------------
    # EXTRACT HIDDEN
    # ---------------
    hiddens = output.get("hiddens") if output else None

    # use every metric passed in as a candidate for callback
    callback_metrics.update(progress_bar_metrics)
    callback_metrics.update(log_metrics)

    # detach all metrics for callbacks to prevent memory leaks
    # no .item() because it will slow things down
    callback_metrics = recursive_detach(callback_metrics)

    return loss, progress_bar_metrics, log_metrics, callback_metrics, hiddens