DataModules: use persistent workers for parallel data loading (#2291)

microsoft · Sep 17, 2024 · d440c56 · d440c56
1 parent 8747aca
commit d440c56
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/torchgeo/datamodules/geo.py b/torchgeo/datamodules/geo.py
@@ -286,6 +286,7 @@ def _dataloader_factory(self, split: str) -> DataLoader[dict[str, Tensor]]:
             batch_sampler=batch_sampler,
             num_workers=self.num_workers,
             collate_fn=self.collate_fn,
+            persistent_workers=self.num_workers > 0,
         )
 
     def train_dataloader(self) -> DataLoader[dict[str, Tensor]]:
@@ -429,6 +430,7 @@ def _dataloader_factory(self, split: str) -> DataLoader[dict[str, Tensor]]:
             shuffle=split == 'train',
             num_workers=self.num_workers,
             collate_fn=self.collate_fn,
+            persistent_workers=self.num_workers > 0,
         )
 
     def train_dataloader(self) -> DataLoader[dict[str, Tensor]]: