Merge branch 'develop' into 'main'

Release 0.2

CHANGELOG.md

@@ -9,6 +9,7 @@ Please add your functional changes to the appropriate section in the PR.
 Keep it human-readable, your future self will thank you!
 
 ## [Unreleased](https://github.com/ecmwf/anemoi-training/compare/0.1.0...HEAD)
+- Make pin_memory of the Dataloader configurable (#64)
 
 ### Added
 

src/anemoi/training/config/dataloader/native_grid.yaml

@@ -1,4 +1,5 @@
 prefetch_factor: 2
+pin_memory: True
 
 num_workers:
   training: 8

src/anemoi/training/data/datamodule.py

@@ -96,6 +96,9 @@ def __init__(self, config: DictConfig) -> None:
             )
             self.config.dataloader.training.end = self.config.dataloader.validation.start - 1
 
+        if not self.config.dataloader.get("pin_memory", True):
+            LOGGER.info("Data loader memory pinning disabled.")
+
     def _check_resolution(self, resolution: str) -> None:
         assert (
             self.config.data.resolution.lower() == resolution.lower()
@@ -185,7 +188,7 @@ def _get_dataloader(self, ds: NativeGridDataset, stage: str) -> DataLoader:
             num_workers=self.config.dataloader.num_workers[stage],
             # use of pinned memory can speed up CPU-to-GPU data transfers
             # see https://pytorch.org/docs/stable/notes/cuda.html#cuda-memory-pinning
-            pin_memory=True,
+            pin_memory=self.config.dataloader.get("pin_memory", True),
             # worker initializer
             worker_init_fn=worker_init_func,
             # prefetch batches