Repository: pytorch
Issue: Support gathering nested lists in DataParallel 
Labels: triaged,module: data parallel

## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->
The [current implementation](https://github.com/pytorch/pytorch/blob/ee628d64b9d2bc8ab26655982aa3a158859a99c0/torch/nn/parallel/scatter_gather.py#L46) of the `gather` function supports simple data structures like tensors and dictionaries but can not handle nested structures such as a list of lists. 
How about supporting them as well?

## Motivation
<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->
Complex networks don't always return simple outputs such as tensors. Consider the following situation where the outputs from the forward pass of a question answering network is a dictionary with keys `loss`, `answers` and `qids`. Notice that `answers` and `qids` are a list of lists. 
```python
# Output from forward pass
 {'loss': tensor(8.6422, device='cuda:0', grad_fn=<ThAddBackward>),
  'answers': [['July 4', 'American Independence'], ['July 1', 'Canadian Independence']],
  'qids': [['C28fbfa5abq#1', 'C28fbfa5abq#2'], ['C2c6ca8b089q#1', 'C2c6ca8b089q#2']}
```

While running this model over multiple GPUs using `nn.DataParallel`, the following is passed as input to the [gather function](https://github.com/pytorch/pytorch/blob/ee628d64b9d2bc8ab26655982aa3a158859a99c0/torch/nn/parallel/scatter_gather.py#L46)

```python
# Input to the `gather` function
[{'loss': tensor(8.6422, device='cuda:0', grad_fn=<ThAddBackward>),
  'answers': [['July 4', 'American Independence'], ['July 1', 'Canadian Independence']],
  'qids': [['C28fbfa5abq#1', 'C28fbfa5abq#2'], ['C2c6ca8b089q#1', 'C2c6ca8b089q#2']},
 {'loss': tensor(8.7826, device='cuda:0', grad_fn=<ThAddBackward>),
  'answers': [['September 16', 'Mexican Independence'], ['November 3', 'Panamanian Independence']],
  'qids': [['C28fbfcdefab#1', 'C28fbfcdefab#2'], ['C2c623ef089q#1', 'C2c623ef089q#2']}]
```
Here's the expected output where the tensors and lists are concatenated together.

```python
# Expected behaviour
{'loss': tensor([8.6422, 8.7826], device='cuda:0', grad_fn=<ThAddBackward>),
  'answers': [['July 4', 'American Independence'], ['July 1', 'Canadian Independence'], ['September 16', 'Mexican Independence'], ['November 3', 'Panamanian Independence']],
  'qids': [['C28fbfa5abq#1', 'C28fbfa5abq#2'], ['C2c6ca8b089q#1', 'C2c6ca8b089q#2'], ['C28fbfcdefab#1', 'C28fbfcdefab#2'], ['C2c623ef089q#1', 'C2c623ef089q#2']]}
```

But this is what happens right now. The `map` clause cannot handle nested lists.

```python
# Current behaviour
{'loss': tensor([8.6422, 8.7826], device='cuda:0', grad_fn=<ThAddBackward>),
  'answers': [['<map object at 0x2ad2445a7f60>', '<map object at 0x2ad2445a75f8>'], ['<map object at 0x2b6862d87358>', '<map object at 0x2b6862d87668>']],
  'qids': [['<map object at 0x2b6862d87668>', '<map object at 0x2b6862d875c0>'], ['<map object at 0x2b6862d87438>', '<map object at 0x2b6862d87828>']]}
```

## Pitch
<!-- A clear and concise description of what you want to happen. -->
One solution is to explicitly handle merging of lists together. This basically boils down to adding an extra condition for handling lists in the `gather` function [here](https://github.com/pytorch/pytorch/blob/ee628d64b9d2bc8ab26655982aa3a158859a99c0/torch/nn/parallel/scatter_gather.py#L62).

```python
 if isinstance(out, list):
     return [item for output in outputs for item in output]
```

## Alternatives
<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->
One workaround is to avoid nested outputs from the forward pass in the first place. This is still a workaround and not a solution.

## Additional context
<!-- Add any other context or screenshots about the feature request here. -->
I'm not sure if this is the right way to handle the larger problem of handling complex data structures.
For example one can always come up with a more nested structure which might not work with the proposed solution either. Should pytorch allow the provision for custom `gather` and `scatter` functions as it has done with `collate_fn`?