Add setting length

openclimatefix · Dec 5, 2023 · 0e3d284 · 0e3d284
1 parent 34b4b30
commit 0e3d284
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/pvnet/data/wind_datamodule.py b/pvnet/data/wind_datamodule.py
@@ -89,11 +89,12 @@ def _get_datapipe(self, start_time, end_time):
         return data_pipeline
 
     def _get_premade_batches_datapipe(self, subdir, shuffle=False):
+        filenames = list(glob.glob(f"{self.batch_dir}/{subdir}/*.nc"))
         data_pipeline = windnet_netcdf_datapipe(
             config_filename=self.configuration,
             keys=["sensor", "nwp"],
-            filenames=list(glob.glob(f"{self.batch_dir}/{subdir}/*.nc")),
-        )
+            filenames=filenames,
+        ).set_length(int(len(filenames) / self.batch_size))
         data_pipeline = (
             data_pipeline.batch(self.batch_size)
             .map(stack_np_examples_into_batch)