diff --git a/IGEV-Stereo/core/stereo_datasets.py b/IGEV-Stereo/core/stereo_datasets.py index c8a0835..d32ca42 100644 --- 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