Update simple_cerebros_random_search.py

david-thrower · Jun 12, 2024 · efce1b6 · efce1b6
1 parent 4c0a279
commit efce1b6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/cerebros/simplecerebrosrandomsearch/simple_cerebros_random_search.py b/cerebros/simplecerebrosrandomsearch/simple_cerebros_random_search.py
@@ -6,7 +6,7 @@
     dense_automl_structural_component \
     import DenseAutoMlStructuralComponent, DenseLateralConnectivity, \
     zero_7_exp_decay, zero_95_exp_decay, simple_sigmoid
-from cerebros.units.units import Unit, InputUnit, FinalDenseUnit
+from cerebros.units.units import Unit, InputUnit, FinalDenseUnit, TernaryDenseLayer
 from cerebros.neuralnetworkfuture.neural_network_future \
     import NeuralNetworkFuture, RealNeuronNeuralNetworkFuture
 # from cmdutil.cmdutil import run_command
@@ -579,7 +579,7 @@ def get_best_model(self):
             tf.keras.models.load_model(
                 self.best_model_path,
                 custom_objects=\
-                    {"TernaryDenseLayer": units.units.TernaryDenseLayer})
+                    {"TernaryDenseLayer": TernaryDenseLayer})
         return best_model
 
 # ->