From 2cdaf422321a6d26880e0fd7a8fca51f0b9ff097 Mon Sep 17 00:00:00 2001 From: Victor Bourgin Date: Fri, 17 Jan 2025 15:41:00 -0800 Subject: [PATCH] Override `num_workers` from 0 to 1 in `get_pytorch_dataloader` Summary: To support a similar interface to the pytorch dataloader, we set `num_workers` to 1 if the input `num_workers` is 0 in `get_pytorch_dataloader` Differential Revision: D68357865 --- src/spdl/dataloader/_pytorch_dataloader.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/spdl/dataloader/_pytorch_dataloader.py b/src/spdl/dataloader/_pytorch_dataloader.py index 2ea74eef..bf11b3e6 100644 --- a/src/spdl/dataloader/_pytorch_dataloader.py +++ b/src/spdl/dataloader/_pytorch_dataloader.py @@ -298,8 +298,13 @@ def get_pytorch_dataloader( if timeout is not None and timeout < 0: raise ValueError(f"`timeout` must be positive. Found: {timeout}.") - if num_workers < 1: + if num_workers < 0: raise ValueError(f"`num_workers` must be greater than 0. Found: {num_workers}") + elif num_workers == 0: + _LG.warning( + "`num_workers` is 0. Setting `num_workers` to 1 for single process dataloading." + ) + num_workers = 1 buffer_size = prefetch_factor * num_workers