| 8 | |
| 9 | |
| 10 | def stack(batch, data_key=None, return_dc=False): |
| 11 | if isinstance(batch[0][data_key], DataContainer): |
| 12 | if batch[0][data_key].stack: |
| 13 | assert isinstance(batch[0][data_key].data, torch.Tensor) |
| 14 | samples = [sample[data_key].data for sample in batch] |
| 15 | return default_collate(samples) |
| 16 | |
| 17 | elif not return_dc: |
| 18 | return [sample[data_key].data for sample in batch] |
| 19 | |
| 20 | else: |
| 21 | return DataContainer([sample[data_key].data for sample in batch]) |
| 22 | |
| 23 | else: |
| 24 | return default_collate([sample[data_key] for sample in batch]) |
| 25 | |
| 26 | |
| 27 | def collate(batch, trans_dict): |