MCPcopy
hub / github.com/PRIME-RL/PRIME / to

Method to

training/verl/protocol.py:224–236  ·  view source on GitHub ↗

move the batch to device Args: device (torch.device, str): torch device Returns: DataProto: the current DataProto

(self, device)

Source from the content-addressed store, hash-verified

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

Callers 15

update_actorMethod · 0.95
update_criticMethod · 0.95
update_actorMethod · 0.95
update_criticMethod · 0.95
offload_fsdp_gradFunction · 0.80
load_fsdp_gradFunction · 0.80
load_fsdp_param_and_gradFunction · 0.80
offload_fsdp_optimizerFunction · 0.80
load_fsdp_optimizerFunction · 0.80

Calls

no outgoing calls

Tested by

no test coverage detected