PinMemory

class torchdata.datapipes.iter.PinMemory(source_datapipe, device=None, pin_memory_fn=<function pin_memory_fn>)

Prefetches one element from the source DataPipe and moves it to pinned memory (functional name: pin_memory). When used with MultiProcessingReadingService, this DataPipe would be kept in the main process to prevent duplicated CUDA context creation.

Parameters:

• source_datapipe – IterDataPipe from which samples are moved to pinned memory.

• device – The device to pin samples.

• pin_memory_fn – Optional callable function to move data to pinned memory. A pin_memory_fn to handle general objects is provided by default.

Example

>>> from torchdata.datapipes.iter import IterableWrapper
>>> dp = IterableWrapper(file_paths).open_files().readlines().map(tokenize_fn).pin_memory()