Struct DistributedDataLoader 

pub struct DistributedDataLoader {
    pub batch_size: usize,
    pub world_size: usize,
    pub rank: usize,
    pub shuffle: bool,
    pub drop_last: bool,
}

Distributed data loader

Fields

batch_size: usize

world_size: usize

rank: usize

shuffle: bool

drop_last: bool

Implementations

impl DistributedDataLoader

pub fn new(batch_size: usize, world_size: usize, rank: usize) -> Self

pub fn shuffle(self, shuffle: bool) -> Self

pub fn drop_last(self, drop: bool) -> Self

pub fn get_batches( &self, data: &Tensor, labels: &Tensor, ) -> Vec<(Tensor, Tensor)>

Get batches for this worker