Generalize loss_function serializer (#74)

openscilab · Jan 20, 2024 · 116416d · 116416d
1 parent 86ce501
commit 116416d
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/pymilo/transporters/lossfunction_transporter.py b/pymilo/transporters/lossfunction_transporter.py
@@ -25,10 +25,10 @@ def serialize(self, data, key, model_type):
         :return: pymilo serialized output of data[key]
         """
         if (
-            (model_type == "SGDClassifier" and key == "loss_function_") or
-            (model_type == "SGDOneClassSVM" and key == "loss_function_") or
-            (model_type == "Perceptron" and key == "loss_function_") or
-            (model_type == "PassiveAggressiveClassifier" and key == "loss_function_")
+            (model_type == "SGDClassifier" and (key == "loss_function_" or key == "_loss_function_")) or
+            (model_type == "SGDOneClassSVM" and (key == "loss_function_" or key == "_loss_function_")) or
+            (model_type == "Perceptron" and (key == "loss_function_" or key == "_loss_function_")) or
+            (model_type == "PassiveAggressiveClassifier" and (key == "loss_function_" or key == "_loss_function_"))
         ):
             data[key] = {
                 "loss": data["loss"]