Update simple_cerebros_random_search.py

david-thrower · Jun 13, 2024 · df7b200 · df7b200
1 parent 7dc1b73
commit df7b200
Showing 1 changed file with 4 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,8 @@
     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, TernaryDenseLayer
+from cerebros.units.units import Unit, InputUnit, FinalDenseUnit
+from custom.custom import TernaryDenseLayer, GPT2Layer
 from cerebros.neuralnetworkfuture.neural_network_future \
     import NeuralNetworkFuture, RealNeuronNeuralNetworkFuture
 # from cmdutil.cmdutil import run_command
@@ -579,7 +580,8 @@ def get_best_model(self):
             tf.keras.models.load_model(
                 self.best_model_path,
                 custom_objects=\
-                    {"TernaryDenseLayer": TernaryDenseLayer})
+                    {"TernaryDenseLayer": TernaryDenseLayer,
+                     "GPT2Layer": GPT2Layer})
         return best_model
 
 # ->