move the batch to device Args: device (torch.device, str): torch device Returns: DataProto: the current DataProto
(self, device)
| 222 | return cls(batch=tensor_dict, non_tensor_batch=non_tensors, meta_info=meta_info) |
| 223 | |
| 224 | def to(self, device) -> 'DataProto': |
| 225 | """move the batch to device |
| 226 | |
| 227 | Args: |
| 228 | device (torch.device, str): torch device |
| 229 | |
| 230 | Returns: |
| 231 | DataProto: the current DataProto |
| 232 | |
| 233 | """ |
| 234 | if self.batch is not None: |
| 235 | self.batch = self.batch.to(device) |
| 236 | return self |
| 237 | |
| 238 | def select(self, batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=None, deepcopy=False) -> 'DataProto': |
| 239 | """Select a subset of the DataProto via batch_keys and meta_info_keys |
no outgoing calls
no test coverage detected