Struct DataLoader 

pub struct DataLoader { /* private fields */ }

Implementations

impl DataLoader

pub fn new(dataset: Dataset, batch_size: usize, shuffle: bool) -> Self

pub fn new_distributed( dataset: Dataset, batch_size: usize, shuffle: bool, rank: usize, world_size: usize, ) -> Self

Create a distributed data loader that shards data across ranks Each rank processes 1/world_size of the data

pub fn reset(&mut self)

pub fn reset_with_seed(&mut self, seed: u64)

Reset with a specific seed (for reproducibility across epochs)

pub fn position(&self) -> usize

Get current position for checkpointing

pub fn set_position(&mut self, pos: usize)

Set position (for resuming from checkpoint)

pub fn num_batches(&self) -> usize

pub fn next_batch( &mut self, device: &Device, ) -> Result<Option<(Tensor, Tensor)>>

pub fn iter<'a>(&'a mut self, device: &'a Device) -> DataLoaderIterator<'a>