Update units.py

david-thrower · Jun 12, 2024 · 8d120be · 8d120be
1 parent dc9de41
commit 8d120be
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/cerebros/units/units.py b/cerebros/units/units.py
@@ -553,8 +553,9 @@ def materialize(self):
             rn_5 = ''
             self.neural_network_layer =\
                 TernaryDenseLayer(
-                    self.n_neurons,
-                    self.activation,
+                    units=self.n_neurons,
+                    input_dim=5,
+                    activation=self.activation,
                     name=f"{self.name}_dns_{rn_5}")(merged_neural_network_layer_input)
             self.materialized = True
         # refactor the lagic below and this class is complete