ray.data.DatasetPipeline.to_torch
ray.data.DatasetPipeline.to_torch#
- DatasetPipeline.to_torch(*, label_column: Optional[str] = None, feature_columns: Optional[Union[List[str], List[List[str]], Dict[str, List[str]]]] = None, label_column_dtype: Optional[torch.dtype] = None, feature_column_dtypes: Optional[Union[torch.dtype, List[torch.dtype], Dict[str, torch.dtype]]] = None, batch_size: int = 1, prefetch_blocks: int = 0, drop_last: bool = False, unsqueeze_label_tensor: bool = True, unsqueeze_feature_tensors: bool = True) torch.utils.data.IterableDataset [source]#
Call
Dataset.to_torch
over the stream of output batches from the pipeline