nip.scenario_base.data.TensorDictDataLoader

class nip.scenario_base.data.TensorDictDataLoader(dataset: TensorDictDataset, **kwargs)

Dataloader class for TensorDict datasets.

Works with a TensorDictDataset subclass. The dataloader will yield tensordicts.

Parameters:

dataset (TensorDictDataset) – The dataset to load.

Methods Summary

__init__(dataset, **kwargs)
_get_iterator()

Attributes

_auto_collation
_index_sampler
multiprocessing_context
dataset
batch_size
num_workers
pin_memory
drop_last
timeout
sampler
pin_memory_device
prefetch_factor

Methods

__init__(dataset: TensorDictDataset, **kwargs)

_get_iterator() → _BaseDataLoaderIter