get_worker_info

→WorkerInfo or None

get_worker_info()

source

Return a WorkerInfo object in a worker process, else None.

Inside the worker_init_fn or dataset __getitem__ / __iter__ of a multi-process DataLoader, this returns the WorkerInfo for the current worker. In the main process — or when num_workers=0 — it returns None.

Returns

WorkerInfo or None

Per-worker context if called from a worker, else None.

Notes

The dominant use case is sharding an IterableDataset across workers: each worker pulls its WorkerInfo.id and advances through a distinct slice of the underlying stream so that samples are not duplicated across the pool.

Examples

>>> class ShardedStream(IterableDataset):
...     def __iter__(self):
...         info = get_worker_info()
...         if info is None:
...             yield from self._all()
...         else:
...             yield from self._slice(info.id, info.num_workers)