Results gathering with varying tensor shapes (e.g. last batch)

[awaelchli]

🐛 Bug

Results object reduction when batch sizes are different won't work because torch.stack get's different input shapes. This can happen if your dataloader returns a smaller batch for the last iteration, for example.

def recursive_stack(result: MutableMapping):
    for k, v in result.items():
        if isinstance(v, dict):
            recursive_stack(v)
        if isinstance(v, list) and len(v) > 0 and isinstance(v[0], torch.Tensor):
            v = torch.stack(v)
            result[k] = v

Context
From slack discussion by @artgor
https://pytorch-lightning.slack.com/archives/CRBLFHY79/p1597604494424600

[awaelchli]

will submit a PR within the next hours