MultiGPU Training. Logging problem
🐛 Bug
When we try logging to the tensorboard loss during last step of epoch with GPU number more than one, we can get exception.
To Reproduce
Start training any model on any dataset with gpu amount over than one, but last batch should contains objects only for the part of GPUs.
I have next error:
ValueError: only one element tensors can be converted to Python scalars
Expected behavior
Success logging to tensorboard.

Environment
OS: Ubuntu
pytorch-lightning==0.7.1
Additional context
I suppose problem in the method (trainer/logging.py):
def reduce_distributed_output(self, output, num_gpus)
# reduce only metrics that have the same number of gpus
elif output[k].size(0) == num_gpus:
    reduced = torch.mean(output[k])
    output[k] = reduced
If we have a last not full batch, we should get mean, isn't it?