Update simple_cerebros_random_search.py

david-thrower · Jun 12, 2024 · 4c0a279 · 4c0a279
1 parent 2b6ccc4
commit 4c0a279
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/cerebros/simplecerebrosrandomsearch/simple_cerebros_random_search.py b/cerebros/simplecerebrosrandomsearch/simple_cerebros_random_search.py
@@ -572,10 +572,14 @@ def run_random_search(self):
 
     def get_best_model(self):
         best_model_path = self.best_model_path
+        # best_model =\
+        #     tf.keras.models.clone_model(
+        #         tf.keras.models.load_model(best_model_path, custom_objects={"TernaryDenseLayer": units.units.TernaryDenseLayer}))
         best_model =\
-            tf.keras.models.clone_model(
-                tf.keras.models.load_model(best_model_path))
-        # best_model = tf.keras.models.load_model(self.best_model_path)
+            tf.keras.models.load_model(
+                self.best_model_path,
+                custom_objects=\
+                    {"TernaryDenseLayer": units.units.TernaryDenseLayer})
         return best_model
 
 # ->