fix steps_per_epoch

knaw-huc · Jan 20, 2025 · 4309f34 · 4309f34
1 parent 50ff44b
commit 4309f34
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/data/manager.py b/src/data/manager.py
@@ -59,8 +59,10 @@ def __init__(self,
 
         # Fill the datasets dictionary with datasets for different partitions
         logging.info("Creating datasets...")
+        steps_per_epoch = config["steps_per_epoch"] \
+            if "steps_per_epoch" in config else None
         self.datasets = self._fill_datasets_dict(file_names, labels,
-                                                 sample_weights, config['steps_per_epoch'])
+                                                 sample_weights, steps_per_epoch)
 
     def _process_raw_data(self) -> Tuple[Dict[str, List[str]],
     Dict[str, List[str]],