From 315b35afae4a2978a7d932940ea0f4a082654abb Mon Sep 17 00:00:00 2001 From: moto <855818+mthrok@users.noreply.github.com> Date: Fri, 17 Jan 2025 09:56:13 -0500 Subject: [PATCH] Add in_order support --- src/spdl/dataloader/_pytorch_dataloader.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/spdl/dataloader/_pytorch_dataloader.py b/src/spdl/dataloader/_pytorch_dataloader.py index eb579ce1..2ea74eef 100644 --- a/src/spdl/dataloader/_pytorch_dataloader.py +++ b/src/spdl/dataloader/_pytorch_dataloader.py @@ -279,6 +279,7 @@ def get_pytorch_dataloader( prefetch_factor: int = 2, persistent_workers: bool = False, pin_memory_device: str | None = None, + in_order: bool = False, ) -> PyTorchDataLoader[U]: from torch.utils.data.dataloader import IterableDataset @@ -336,4 +337,5 @@ def get_pytorch_dataloader( num_workers=num_workers, timeout=timeout, buffer_size=buffer_size, + output_order="input" if in_order else "completion", )