# Buggy code
def log_train_epoch_end_metrics(
    self,
    epoch_output,
    checkpoint_accumulator,
    early_stopping_accumulator,
    num_optimizers,
):
    # epoch output is a list. Each item in that list has all the outputs per optimizer
    # epoch_output[optimizer_idx][training_step_idx][tbptt_index]
    # remember that not using truncated backprop is equivalent with truncated back prop of len(1)

    model = self.trainer.get_model()

    epoch_callback_metrics = {}

    # -----------------------
    # Calculate epoch callback values if given
    # -----------------------
    if checkpoint_accumulator.num_values > 0:
        epoch_callback_metrics["checkpoint_on"] = checkpoint_accumulator.mean()

    if early_stopping_accumulator.num_values > 0:
        epoch_callback_metrics["early_stop_on"] = early_stopping_accumulator.mean()

    # ------------------------
    # determine if using a result obj
    # ------------------------
    # [optimizer_idx][training_step_idx][tbptt_index]
    opt_idx_outputs = epoch_output[0]

    # TODO: deprecate 1.0
    try:
        sample_obj = (
            opt_idx_outputs[0][0]
            if isinstance(opt_idx_outputs[0], list)
            else opt_idx_outputs[0]
        )
        is_result_obj = len(epoch_output) > 0 and isinstance(sample_obj, Result)
        is_1_0_result = is_result_obj and "extra" in sample_obj
    except IndexError as e:
        is_result_obj = False
        is_1_0_result = False

    # ------------------
    # NEW 1.0.0 PATH
    # ------------------
    if is_1_0_result:
        # lightning module hook
        self.training_epoch_end(model, epoch_output, num_optimizers)

        # log/aggregate metrics automatically
        epoch_log_metrics, epoch_progress_bar_metrics = (
            self.__auto_reduce_results_on_epoch_end(epoch_output)
        )

    # TODO: deprecate 1.0
    else:
        out = self.__run_legacy_training_epoch_end(
            num_optimizers, epoch_output, model, is_result_obj, epoch_callback_metrics
        )
        epoch_log_metrics, epoch_progress_bar_metrics, epoch_callback_metrics = out

    # it will perform reduction over epoch and return log metrics
    cached_epoch_log_metrics = self.cached_results.get_epoch_log_metrics()
    cached_epoch_pbar_metrics = self.cached_results.get_epoch_pbar_metrics()

    # update
    epoch_log_metrics.update(cached_epoch_log_metrics)
    epoch_progress_bar_metrics.update(cached_epoch_pbar_metrics)

    # --------------------------
    # track results
    # --------------------------
    # add the metrics to the loggers and callbacks
    if epoch_log_metrics and len(epoch_log_metrics) > 0:
        self.log_metrics(epoch_log_metrics, {})
        self.callback_metrics.update(epoch_log_metrics)

    # add metrics to callbacks
    self.callback_metrics.update(epoch_callback_metrics)

    # add metrics to progress_bar and callbacks
    if len(epoch_progress_bar_metrics) > 0:
        self.add_progress_bar_metrics(epoch_progress_bar_metrics)
        self.callback_metrics.update(epoch_progress_bar_metrics)

    # reset epoch loop result for next epoch
    self.cached_results.reset()


# Traceback (RuntimeError)
GPU available: True, used: True
TPU available: False, using: 0 TPU cores
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]

| Name  | Type   | Params
---------------------------------
0 | layer | Linear | 66
/home/user/.conda/envs/env/lib/python3.8/site-packages/pytorch_lightning/utilities/distributed.py:45: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 104 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.
warnings.warn(*args, **kwargs)
Validation sanity check: 0it [00:00, ?it/s]/home/user/.conda/envs/env/lib/python3.8/site-packages/torch/nn/parallel/_functions.py:61: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.
warnings.warn('Was asked to gather along dimension 0, but all '
/home/user/.conda/envs/env/lib/python3.8/site-packages/pytorch_lightning/utilities/distributed.py:45: UserWarning: The dataloader, train dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 104 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.
warnings.warn(*args, **kwargs)
Epoch 1:  50%|████████████████Traceback (most recent call last):████████████████████████████████████                                                                                      | 4/8 [00:00<00:00, 184.41it/s, loss=0.497, v_num=53]
File "dp_bug.py", line 118, in <module>
main()
File "dp_bug.py", line 114, in main
trainer.fit(model)
File "/home/user/.conda/envs/env/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 440, in fit
results = self.accelerator_backend.train()
File "/home/user/.conda/envs/env/lib/python3.8/site-packages/pytorch_lightning/accelerators/dp_accelerator.py", line 97, in train
results = self.train_or_test()
File "/home/user/.conda/envs/env/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 53, in train_or_test
results = self.trainer.train()
File "/home/user/.conda/envs/env/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 483, in train
self.train_loop.run_training_epoch()
File "/home/user/.conda/envs/env/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 557, in run_training_epoch
self.trainer.run_evaluation(test_mode=False)
File "/home/user/.conda/envs/env/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 599, in run_evaluation
eval_loop_results = self.evaluation_loop.log_epoch_metrics(deprecated_eval_results, epoch_logs, test_mode)
File "/home/user/.conda/envs/env/lib/python3.8/site-packages/pytorch_lightning/trainer/evaluation_loop.py", line 210, in log_epoch_metrics
eval_loop_results = self.trainer.logger_connector.on_evaluation_epoch_end(
File "/home/user/.conda/envs/env/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/logger_connector.py", line 113, in on_evaluation_epoch_end
self._log_on_evaluation_epoch_end_metrics(epoch_logs)
File "/home/user/.conda/envs/env/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/logger_connector.py", line 178, in _log_on_evaluation_epoch_end_metrics
reduced_epoch_metrics = dl_metrics[0].__class__.reduce_on_epoch_end(dl_metrics)
File "/home/user/.conda/envs/env/lib/python3.8/site-packages/pytorch_lightning/core/step_result.py", line 433, in reduce_on_epoch_end
recursive_stack(result)
File "/home/user/.conda/envs/env/lib/python3.8/site-packages/pytorch_lightning/core/step_result.py", line 552, in recursive_stack
result[k] = collate_tensors(v)
File "/home/user/.conda/envs/env/lib/python3.8/site-packages/pytorch_lightning/core/step_result.py", line 574, in collate_tensors
return torch.stack(items)
RuntimeError: All input tensors must be on the same device. Received cuda:3 and cuda:1
Exception ignored in: <function tqdm.__del__ at 0x7fcf54050a60>
Traceback (most recent call last):
File "/home/user/.conda/envs/env/lib/python3.8/site-packages/tqdm/std.py", line 1087, in __del__
File "/home/user/.conda/envs/env/lib/python3.8/site-packages/tqdm/std.py", line 1294, in close
File "/home/user/.conda/envs/env/lib/python3.8/site-packages/tqdm/std.py", line 1472, in display
File "/home/user/.conda/envs/env/lib/python3.8/site-packages/tqdm/std.py", line 1090, in __repr__
File "/home/user/.conda/envs/env/lib/python3.8/site-packages/tqdm/std.py", line 1434, in format_dict
TypeError: cannot unpack non-iterable NoneType object

# Fixed code
def log_train_epoch_end_metrics(
    self,
    epoch_output,
    checkpoint_accumulator,
    early_stopping_accumulator,
    num_optimizers,
):
    # epoch output is a list. Each item in that list has all the outputs per optimizer
    # epoch_output[optimizer_idx][training_step_idx][tbptt_index]
    # remember that not using truncated backprop is equivalent with truncated back prop of len(1)

    model = self.trainer.get_model()

    epoch_callback_metrics = {}

    # -----------------------
    # Calculate epoch callback values if given
    # -----------------------
    if checkpoint_accumulator.num_values > 0:
        epoch_callback_metrics["checkpoint_on"] = checkpoint_accumulator.mean()

    if early_stopping_accumulator.num_values > 0:
        epoch_callback_metrics["early_stop_on"] = early_stopping_accumulator.mean()

    # ------------------------
    # determine if using a result obj
    # ------------------------
    # [optimizer_idx][training_step_idx][tbptt_index]
    opt_idx_outputs = epoch_output[0]

    # TODO: deprecate 1.0
    try:
        sample_obj = (
            opt_idx_outputs[0][0]
            if isinstance(opt_idx_outputs[0], list)
            else opt_idx_outputs[0]
        )
        is_result_obj = len(epoch_output) > 0 and isinstance(sample_obj, Result)
        is_1_0_result = is_result_obj and "extra" in sample_obj
    except IndexError as e:
        is_result_obj = False
        is_1_0_result = False

    # ------------------
    # NEW 1.0.0 PATH
    # ------------------
    if is_1_0_result:
        # lightning module hook
        self.training_epoch_end(model, epoch_output, num_optimizers)

        # log/aggregate metrics automatically
        epoch_log_metrics, epoch_progress_bar_metrics = (
            self.__auto_reduce_results_on_epoch_end(epoch_output)
        )

    # TODO: deprecate 1.0
    else:
        out = self.__run_legacy_training_epoch_end(
            num_optimizers, epoch_output, model, is_result_obj, epoch_callback_metrics
        )
        epoch_log_metrics, epoch_progress_bar_metrics, epoch_callback_metrics = out

    # it will perform reduction over epoch and return log metrics
    cached_epoch_log_metrics = self.cached_results.get_epoch_log_metrics()
    cached_epoch_pbar_metrics = self.cached_results.get_epoch_pbar_metrics()

    # update
    epoch_log_metrics.update(cached_epoch_log_metrics)
    epoch_progress_bar_metrics.update(cached_epoch_pbar_metrics)

    # --------------------------
    # track results
    # --------------------------
    # add the metrics to the loggers and callbacks
    if epoch_log_metrics and len(epoch_log_metrics) > 0:
        self.log_metrics(epoch_log_metrics, {})
        self.callback_metrics.update(epoch_log_metrics)

    # add metrics to callbacks
    self.callback_metrics.update(epoch_callback_metrics)

    # add metrics to progress_bar and callbacks
    if len(epoch_progress_bar_metrics) > 0:
        self.add_progress_bar_metrics(epoch_progress_bar_metrics)
        self.callback_metrics.update(epoch_progress_bar_metrics)

    # reset epoch loop result for next epoch
    self.cached_results.reset()