support keras2 & keras3

cta-observatory · Oct 18, 2024 · a953e4e · a953e4e
1 parent 43547b2
commit a953e4e
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/dl1_data_handler/loader.py b/dl1_data_handler/loader.py
@@ -13,7 +13,9 @@ def __init__(
         tasks,
         batch_size=64,
         random_seed=0,
+        **kwargs,
     ):
+        super().__init__(**kwargs)
         "Initialization"
         self.DLDataReader = DLDataReader
         self.indices = indices
@@ -82,4 +84,7 @@ def __getitem__(self, index):
                 ),
                 axis=1,
             )
+        # Temp fix for supporting keras2 & keras3
+        if int(keras.__version__.split(".")[0]) >= 3:
+            features = features["input"]
         return features, labels