| 1048 | return loss, progress_bar_metrics, log_metrics, callback_metrics, hiddens |
| 1049 | |
| 1050 | def reduce_distributed_output(self, output, num_gpus): |
| 1051 | if num_gpus <= 1: |
| 1052 | return output |
| 1053 | |
| 1054 | # when using DP, we get one output per gpu |
| 1055 | # average outputs and return |
| 1056 | if type(output) is torch.Tensor: |
| 1057 | return output.mean() |
| 1058 | |
| 1059 | for k, v in output.items(): |
| 1060 | # recurse on nested dics |
| 1061 | if isinstance(output[k], dict): |
| 1062 | output[k] = self.reduce_distributed_output(output[k], num_gpus) |
| 1063 | |
| 1064 | # do nothing when there's a scalar |
| 1065 | elif isinstance(output[k], torch.Tensor) and output[k].dim() == 0: |
| 1066 | pass |
| 1067 | |
| 1068 | # reduce only metrics that have the same number of gpus |
| 1069 | elif output[k].size(0) == num_gpus: |
| 1070 | reduced = torch.mean(output[k]) |
| 1071 | output[k] = reduced |
| 1072 | return output |
| 1073 | |
| 1074 | def clip_gradients(self): |
| 1075 | if self.gradient_clip_val > 0: |