# Buggy code
def _track_callback_metrics(self, eval_results, using_eval_result):
    if len(eval_results) > 0 and (
        eval_results[0] is None or not isinstance(eval_results[0], Result)
    ):
        return

    if using_eval_result:
        if isinstance(eval_results, list):
            for eval_result in eval_results:
                self.trainer.logger_connector.callback_metrics.update(
                    eval_result.callback_metrics
                )
                self.trainer.logger_connector.evaluation_callback_metrics.update(
                    eval_result.callback_metrics
                )
        else:
            self.trainer.logger_connector.callback_metrics.update(
                eval_results.callback_metrics
            )
            self.trainer.logger_connector.evaluation_callback_metrics.update(
                eval_results.callback_metrics
            )
    else:
        flat = {}
        if isinstance(eval_results, list):
            for eval_result in eval_results:
                # with a scalar return, auto set it to "val_loss" for callbacks
                if isinstance(eval_result, torch.Tensor):
                    flat = {"val_loss": eval_result}
                elif isinstance(eval_result, dict):
                    flat = flatten_dict(eval_result)

                # removing val_loss magic word to map to checkpoint + ES callback
                if "val_loss" in flat:
                    flat["checkpoint_on"] = flat["val_loss"]
                    flat["early_stop_on"] = flat["val_loss"]
                self.trainer.logger_connector.callback_metrics.update(flat)
                self.trainer.logger_connector.evaluation_callback_metrics.update(flat)
        else:
            # with a scalar return, auto set it to "val_loss" for callbacks
            if isinstance(eval_results, torch.Tensor):
                flat = {"val_loss": eval_results}
            else:
                flat = flatten_dict(eval_results)

            # removing val_loss magic word to map to checkpoint + ES callback
            if "val_loss" in flat:
                flat["checkpoint_on"] = flat["val_loss"]
                flat["early_stop_on"] = flat["val_loss"]
            self.trainer.logger_connector.callback_metrics.update(flat)
            self.trainer.logger_connector.evaluation_callback_metrics.update(flat)


# 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 _track_callback_metrics(self, eval_results, using_eval_result):
    if len(eval_results) > 0 and (
        eval_results[0] is None or not isinstance(eval_results[0], Result)
    ):
        return

    if using_eval_result:
        if isinstance(eval_results, list):
            for eval_result in eval_results:
                self.trainer.logger_connector.callback_metrics.update(
                    eval_result.callback_metrics
                )
                self.trainer.logger_connector.evaluation_callback_metrics.update(
                    eval_result.callback_metrics
                )
        else:
            self.trainer.logger_connector.callback_metrics.update(
                eval_results.callback_metrics
            )
            self.trainer.logger_connector.evaluation_callback_metrics.update(
                eval_results.callback_metrics
            )
    else:
        flat = {}
        if isinstance(eval_results, list):
            for eval_result in eval_results:
                # with a scalar return, auto set it to "val_loss" for callbacks
                if isinstance(eval_result, torch.Tensor):
                    flat = {"val_loss": eval_result}
                elif isinstance(eval_result, dict):
                    flat = flatten_dict(eval_result)

                # removing val_loss magic word to map to checkpoint + ES callback
                if "val_loss" in flat:
                    flat["checkpoint_on"] = flat["val_loss"]
                    flat["early_stop_on"] = flat["val_loss"]
                self.trainer.logger_connector.callback_metrics.update(flat)
                self.trainer.logger_connector.evaluation_callback_metrics.update(flat)
        else:
            # with a scalar return, auto set it to "val_loss" for callbacks
            if isinstance(eval_results, torch.Tensor):
                flat = {"val_loss": eval_results}
            else:
                flat = flatten_dict(eval_results)

            # removing val_loss magic word to map to checkpoint + ES callback
            if "val_loss" in flat:
                flat["checkpoint_on"] = flat["val_loss"]
                flat["early_stop_on"] = flat["val_loss"]
            self.trainer.logger_connector.callback_metrics.update(flat)
            self.trainer.logger_connector.evaluation_callback_metrics.update(flat)