set 12 threads to prepare datasets

2023-05-03 04:14:05 +08:00 · 2023-05-03 04:14:05 +08:00 · 0404f5b5b0
commit 0404f5b5b0
parent 3b68318ed5
1 changed files with 1 additions and 1 deletions
--- a/IGEV-Stereo/core/stereo_datasets.py
+++ b/IGEV-Stereo/core/stereo_datasets.py
@ -353,7 +353,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