diff --git a/src/model.py b/src/model.py
index e113fb0..39f7ab4 100644
--- a/src/model.py
+++ b/src/model.py
@@ -34,7 +34,7 @@ def build_model(hyperparams) -> Model:
                 units=64, activation="relu", kernel_regularizer=l2(hyperparams.l2_rate)
             ),
             Dropout(rate=0.3),
-            Dense(units=32, activation="softmax"),
+            Dense(units=len(BASES), activation="softmax"),
         ]
     )
     model.compile(