| 30 | (-1 means the CPU), with dictionary support. |
| 31 | """ |
| 32 | def gather_map(outputs): |
| 33 | out = outputs[0] |
| 34 | if torch.is_tensor(out): |
| 35 | # MJY(20180330) HACK:: force nr_dims > 0 |
| 36 | if out.dim() == 0: |
| 37 | outputs = [o.unsqueeze(0) for o in outputs] |
| 38 | return Gather.apply(target_device, dim, *outputs) |
| 39 | elif out is None: |
| 40 | return None |
| 41 | elif isinstance(out, collections.Mapping): |
| 42 | return {k: gather_map([o[k] for o in outputs]) for k in out} |
| 43 | elif isinstance(out, collections.Sequence): |
| 44 | return type(out)(map(gather_map, zip(*outputs))) |
| 45 | return gather_map(outputs) |
| 46 | |
| 47 | |