Method to

common/batch.py:19–23

(self, device: torch.device, non_blocking: bool = False)

    raise NotImplementedError

  def to(self, device: torch.device, non_blocking: bool = False):
    args = {}
    for feature_name, feature_value in self.as_dict().items():
      args[feature_name] = feature_value.to(device=device, non_blocking=non_blocking)
    return self.__class__(**args)

  def record_stream(self, stream: torch.cuda.streams.Stream) -> None:
    for feature_value in self.as_dict().values():