Error in returning Dict from training_step with multiple GPUs

[kchuang625]

🐛 Bug

When using multiple GPUs with 'dp', the error RuntimeError: grad can be implicitly created only for scalar outputs occurs if I utilized training_step function like this:

def training_step(self, batch, batch_idx):
    ...
    return {'loss': loss}

Please reproduce using the BoringModel

https://colab.research.google.com/drive/1hmHqYHPOqDlZUAF7-9zcCvobrvSPt7W5?usp=sharing

Expected behavior

It is supposed to work fine to return Dict with loss key.

A quick solution

Return loss tensor directly from training_step function:

def training_step(self, batch, batch_idx):
    ...
    return loss

Environment

• PyTorchLightning Version: 1.2.0
• PyTorch Version: 1.7.0
• OS: Linux
• Python version: 3.8
• CUDA/cuDNN version: 10.2

cc. @carmocca

[edenlightning]

Need to support reducing returned dict and not just tensors.