fix prefetch factor

openclimatefix · Nov 21, 2023 · d4b7fe4 · d4b7fe4
1 parent 4e1cd7b
commit d4b7fe4
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 7 deletions.
diff --git a/pvnet/data/datamodule.py b/pvnet/data/datamodule.py
@@ -3,14 +3,16 @@
 
 import numpy as np
 import torch
+from torch.utils.data import DataLoader
+from torch.utils.data.datapipes.datapipe import IterDataPipe
+from torch.utils.data.datapipes.iter import FileLister
+from torch.utils.data.datapipes._decorator import functional_datapipe
+
 from lightning.pytorch import LightningDataModule
 from ocf_datapipes.training.pvnet import pvnet_datapipe
 from ocf_datapipes.utils.consts import BatchKey
 from ocf_datapipes.utils.utils import stack_np_examples_into_batch
-from torch.utils.data import DataLoader
-from torch.utils.data.datapipes._decorator import functional_datapipe
-from torch.utils.data.datapipes.datapipe import IterDataPipe
-from torch.utils.data.datapipes.iter import FileLister
+
 
 
 def copy_batch_to_device(batch, device):
@@ -70,7 +72,7 @@ def __init__(
         configuration=None,
         batch_size=16,
         num_workers=0,
-        prefetch_factor=2,
+        prefetch_factor=None,
         train_period=[None, None],
         val_period=[None, None],
         test_period=[None, None],

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -96,7 +96,7 @@ def sample_datamodule():
         configuration=None,
         batch_size=2,
         num_workers=0,
-        prefetch_factor=2,
+        prefetch_factor=None,
         train_period=[None, None],
         val_period=[None, None],
         test_period=[None, None],

diff --git a/tests/data/test_datamodule.py b/tests/data/test_datamodule.py
@@ -6,7 +6,7 @@ def test_init():
         configuration=None,
         batch_size=2,
         num_workers=0,
-        prefetch_factor=2,
+        prefetch_factor=Mone,
         train_period=[None, None],
         val_period=[None, None],
         test_period=[None, None],