# Buggy code
def cache_result(self) -> None:
    """
    This function is called after every hook
    and store the result object
    """
    with self.trainer.profiler.profile("cache_result"):
        model_ref = self.trainer.get_model()

        # extract hook results
        hook_result = model_ref._results

        if len(hook_result) == 1:
            model_ref._current_hook_fx_name = None
            model_ref._current_fx_name = ""
            return

        # extract model information
        fx_name, dataloader_idx = self.current_model_info()

        # add only if anything as been logged
        # default len is 1 due to _internals

        if fx_name not in self._internals:
            self._internals[fx_name] = HookResultStore(fx_name)

        extra_info = {}
        if self.has_split_and_opt_idx:
            extra_info = self.extra_info

        # attach capture batch_size
        Result.attach_batch_size(self._batch_size, hook_result)

        hook_result.detach()
        if self.trainer.move_metrics_to_cpu:
            hook_result.cpu()

        self._internals[fx_name].append(
            hook_result, dataloader_idx=dataloader_idx, extra_info=extra_info
        )

        # update logged_metrics, progress_bar_metrics, callback_metrics
        self.update_logger_connector(fx_name)

        self.reset_model()


# 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 cache_result(self) -> None:
    """
    This function is called after every hook
    and store the result object
    """
    with self.trainer.profiler.profile("cache_result"):
        model_ref = self.trainer.get_model()

        # extract hook results
        hook_result = model_ref._results

        if len(hook_result) == 1:
            model_ref._current_hook_fx_name = None
            model_ref._current_fx_name = ""
            return

        # extract model information
        fx_name, dataloader_idx = self.current_model_info()

        # add only if anything as been logged
        # default len is 1 due to _internals

        if fx_name not in self._internals:
            self._internals[fx_name] = HookResultStore(fx_name)

        extra_info = {}
        if self.has_split_and_opt_idx:
            extra_info = self.extra_info

        # attach capture batch_size
        Result.attach_batch_size(self._batch_size, hook_result)

        hook_result.detach()
        if self.trainer.move_metrics_to_cpu:
            hook_result.cpu()
        elif self.trainer.use_dp:
            hook_result.to(torch.device("cuda", self.trainer.root_gpu))

        self._internals[fx_name].append(
            hook_result, dataloader_idx=dataloader_idx, extra_info=extra_info
        )

        # update logged_metrics, progress_bar_metrics, callback_metrics
        self.update_logger_connector(fx_name)

        self.reset_model()