Update units.py

david-thrower · Jun 12, 2024 · 92dc2f1 · 92dc2f1
1 parent 7b58698
commit 92dc2f1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/cerebros/units/units.py b/cerebros/units/units.py
@@ -554,7 +554,7 @@ def materialize(self):
             self.neural_network_layer =\
                 TernaryDenseLayer(
                     units=self.n_neurons,
-                    input_dim=5,
+                    input_dim=merged_neural_network_layer_input.shape[-1],
                     # activation=self.activation,
                     name=f"{self.name}_dns_{rn_5}")(merged_neural_network_layer_input)
             self.materialized = True