Repository: Lightning-AI/pytorch-lightning

Issue:
Results gathering with varying tensor shapes (e.g. last batch)
<!-- 
### Common bugs:
1. Tensorboard not showing in Jupyter-notebook see [issue 79](https://github.com/PyTorchLightning/pytorch-lightning/issues/79).    
2. PyTorch 1.1.0 vs 1.2.0 support [see FAQ](https://github.com/PyTorchLightning/pytorch-lightning#faq)    
-->

## 🐛 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.

```python
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



Discussion:
will submit a PR within the next hours

Resolving patch:
diff --git a/pytorch_lightning/core/step_result.py b/pytorch_lightning/core/step_result.py
--- a/pytorch_lightning/core/step_result.py
+++ b/pytorch_lightning/core/step_result.py
@@ -1,6 +1,6 @@
 import numbers
 from copy import copy
-from typing import Optional, Dict, Union, Sequence, Callable, MutableMapping, Any
+from typing import Optional, Dict, Union, Sequence, Callable, MutableMapping, Any, List, Tuple
 
 import torch
 from torch import Tensor
@@ -405,19 +405,23 @@ def recursive_stack(result: MutableMapping):
         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
+        result[k] = collate_tensors(v)
 
 
-def recursive_padded_stack(result: MutableMapping):
-    for k, v in result.items():
-        if isinstance(v, dict):
-            recursive_stack(v)
+def collate_tensors(items: Union[List, Tuple]) -> Union[Tensor, List, Tuple]:
+    if not items or not isinstance(items, (list, tuple)) or any(not isinstance(item, Tensor) for item in items):
+        # items is not a sequence, empty, or contains non-tensors
+        return items
+
+    if all(item.ndim == 0 for item in items):
+        # all tensors are scalars, we need to stack
+        return torch.stack(items)
+
+    if all(item.ndim >= 1 and item.shape[1:] == items[0].shape[1:] for item in items):
+        # we can concatenate along the first dimension
+        return torch.cat(items)
 
-        if isinstance(v, list) and len(v) > 0 and isinstance(v[0], torch.Tensor):
-            v = torch.stack(v)
-            result[k] = v
+    return items
 
 
 class TrainResult(Result):


Test patch:
diff --git a/tests/core/test_results.py b/tests/core/test_results.py
--- a/tests/core/test_results.py
+++ b/tests/core/test_results.py
@@ -174,3 +174,65 @@ def test_result_obj_predictions_ddp_spawn(tmpdir):
         predictions = torch.load(prediction_file)
         size += len(predictions)
     assert size == len(dm.mnist_test)
+
+
+def test_result_gather_stack():
+    """ Test that tensors get concatenated when they all have the same shape. """
+    outputs = [
+        {"foo": torch.zeros(4, 5)},
+        {"foo": torch.zeros(4, 5)},
+        {"foo": torch.zeros(4, 5)},
+    ]
+    result = Result.gather(outputs)
+    assert isinstance(result["foo"], torch.Tensor)
+    assert list(result["foo"].shape) == [12, 5]
+
+
+def test_result_gather_concatenate():
+    """ Test that tensors get concatenated when they have varying size in first dimension. """
+    outputs = [
+        {"foo": torch.zeros(4, 5)},
+        {"foo": torch.zeros(8, 5)},
+        {"foo": torch.zeros(3, 5)},
+    ]
+    result = Result.gather(outputs)
+    assert isinstance(result["foo"], torch.Tensor)
+    assert list(result["foo"].shape) == [15, 5]
+
+
+def test_result_gather_scalar():
+    """ Test that 0-dim tensors get gathered and stacked correctly. """
+    outputs = [
+        {"foo": torch.tensor(1)},
+        {"foo": torch.tensor(2)},
+        {"foo": torch.tensor(3)},
+    ]
+    result = Result.gather(outputs)
+    assert isinstance(result["foo"], torch.Tensor)
+    assert list(result["foo"].shape) == [3]
+
+
+def test_result_gather_different_shapes():
+    """ Test that tensors of varying shape get gathered into a list. """
+    outputs = [
+        {"foo": torch.tensor(1)},
+        {"foo": torch.zeros(2, 3)},
+        {"foo": torch.zeros(1, 2, 3)},
+    ]
+    result = Result.gather(outputs)
+    expected = [torch.tensor(1), torch.zeros(2, 3), torch.zeros(1, 2, 3)]
+    assert isinstance(result["foo"], list)
+    assert all(torch.eq(r, e).all() for r, e in zip(result["foo"], expected))
+
+
+def test_result_gather_mixed_types():
+    """ Test that a collection of mixed types gets gathered into a list. """
+    outputs = [
+        {"foo": 1.2},
+        {"foo": ["bar", None]},
+        {"foo": torch.tensor(1)},
+    ]
+    result = Result.gather(outputs)
+    expected = [1.2, ["bar", None], torch.tensor(1)]
+    assert isinstance(result["foo"], list)
+    assert result["foo"] == expected