diff --git a/src/dataloaders/build_dataloader.py b/src/dataloaders/build_dataloader.py index af229d3..b780eca 100644 --- a/src/dataloaders/build_dataloader.py +++ b/src/dataloaders/build_dataloader.py @@ -44,8 +44,10 @@ def build_torch_dataloader(self, torch_dataset): torch_dataset, batch_size=1, # batched inference/eval not implemented shuffle=False, + num_workers=self.args.num_workers, pin_memory=torch.cuda.is_available(), collate_fn=self.collate_fn, + persistent_workers=(self.args.num_workers > 0), ) return torch_data_loader