Unfold the first n dims as new batch dim
(data: 'DataProto', batch_dims=2)
| 130 | |
| 131 | |
| 132 | def unfold_batch_dim(data: 'DataProto', batch_dims=2): |
| 133 | """ |
| 134 | Unfold the first n dims as new batch dim |
| 135 | """ |
| 136 | tensor: TensorDict = data.batch |
| 137 | non_tensor = data.non_tensor_batch |
| 138 | tensor.auto_batch_size_(batch_dims=batch_dims) |
| 139 | tensor = tensor.view(-1) |
| 140 | |
| 141 | batch_size = tensor.batch_size[0] |
| 142 | |
| 143 | non_tensor_new = {} |
| 144 | |
| 145 | for key, val in non_tensor.items(): |
| 146 | non_tensor_new[key] = np.reshape(val, newshape=(batch_size, *val.shape[batch_dims:])) |
| 147 | |
| 148 | return DataProto(batch=tensor, non_tensor_batch=non_tensor_new, meta_info=data.meta_info) |
| 149 | |
| 150 | |
| 151 | def collate_fn(x: list['DataProtoItem']): |