From 46cc5b9c5b22672e819ad63ad22bd47762cc9f12 Mon Sep 17 00:00:00 2001
From: Daoyuan Li <94409450+DaoyuanLi2816@users.noreply.github.com>
Date: Tue, 30 Jun 2026 22:17:16 -0700
Subject: [PATCH] Fix DataLoader crash when data.num_workers=0

_build_train_dataloader always passed persistent_workers=True and
prefetch_factor=4. torch's DataLoader rejects both with a ValueError when
num_workers=0, which is a valid setting for single-GPU or debugging runs.
Only pass those two options when num_workers > 0.
---
 deepspec/trainer/base_trainer.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/deepspec/trainer/base_trainer.py b/deepspec/trainer/base_trainer.py
index 6b2eef1..4134ab3 100644
--- a/deepspec/trainer/base_trainer.py
+++ b/deepspec/trainer/base_trainer.py
@@ -294,17 +294,21 @@ def _build_train_dataloader(self, start_offset_samples=0, num_samples=None):
             start_global_offset_samples=start_offset_samples,
             num_samples=num_samples,
         )
-        return DataLoader(
-            self.train_dataset,
+        num_workers = int(self.args.data.num_workers)
+        dataloader_kwargs = dict(
             batch_size=int(self.args.train.local_batch_size),
             sampler=sampler,
             collate_fn=self.data_collator_cls(),
-            num_workers=int(self.args.data.num_workers),
+            num_workers=num_workers,
             pin_memory=True,
             drop_last=True,
-            persistent_workers=True,
-            prefetch_factor=4,
         )
+        # prefetch_factor / persistent_workers are only valid with worker
+        # processes; passing them when num_workers=0 raises a ValueError.
+        if num_workers > 0:
+            dataloader_kwargs["persistent_workers"] = True
+            dataloader_kwargs["prefetch_factor"] = 4
+        return DataLoader(self.train_dataset, **dataloader_kwargs)
 
     def run_batch(self, batch):
         raise NotImplementedError