minor bugfix in test KerasSequenceCNN

uio-bmi · Apr 18, 2024 · 1bea581 · 1bea581
1 parent 1a1d8c4
commit 1bea581
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/immuneML/ml_methods/BinaryFeatureClassifier.py b/immuneML/ml_methods/BinaryFeatureClassifier.py
@@ -336,7 +336,7 @@ def store(self, path: Path, feature_names=None, details_path: Path = None):
         del custom_vars["result_path"]
 
         if self.label:
-            custom_vars["label"] = {key.lstrip("_"): value for key, value in vars(self.label).items()}
+            custom_vars["label"] = self.label.get_desc_for_storage()
 
         params_path = path / "custom_params.yaml"
         with params_path.open('w') as file:

diff --git a/immuneML/ml_methods/KerasSequenceCNN.py b/immuneML/ml_methods/KerasSequenceCNN.py
@@ -200,7 +200,7 @@ def fit_by_cross_validation(self, encoded_data: EncodedData, label: Label = None
     def store(self, path: Path, feature_names=None, details_path: Path = None):
         PathBuilder.build(path)
 
-        self.model.save(path / "model")
+        self.model.save(path / "model.keras")
 
         custom_vars = copy.deepcopy(vars(self))
         del custom_vars["model"]
@@ -228,7 +228,7 @@ def load(self, path):
                 else:
                     setattr(self, param, value)
 
-        self.model = keras.models.load_model(path / "model")
+        self.model = keras.models.load_model(path / "model.keras")
 
     def check_if_exists(self, path):
         return self.model is not None