From c83b93e0bfcd8554785a3d233e5d15d7f9666b47 Mon Sep 17 00:00:00 2001
From: HTensor <27774381+HTensor@users.noreply.github.com>
Date: Wed, 3 May 2023 04:14:05 +0800
Subject: [PATCH] set 12 threads to prepare datasets

---
 IGEV-Stereo/core/stereo_datasets.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/IGEV-Stereo/core/stereo_datasets.py b/IGEV-Stereo/core/stereo_datasets.py
index 20a4e83..d071de0 100644
--- a/IGEV-Stereo/core/stereo_datasets.py
+++ b/IGEV-Stereo/core/stereo_datasets.py
@@ -352,7 +352,7 @@ def fetch_dataloader(args):
         train_dataset = new_dataset if train_dataset is None else train_dataset + new_dataset
 
     train_loader = data.DataLoader(train_dataset, batch_size=args.batch_size, 
-        pin_memory=True, shuffle=True, num_workers=int(os.environ.get('SLURM_CPUS_PER_TASK', 6))-2, drop_last=True)
+        pin_memory=True, shuffle=True, num_workers=12, drop_last=True)
 
     logging.info('Training with %d image pairs' % len(train_dataset))
     return train_loader