diff --git a/src/api.py b/src/api.py
index 45a9319..1f069b5 100644
--- a/src/api.py
+++ b/src/api.py
@@ -1,5 +1,6 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
+
 from model import infer_sequence
 
 app = FastAPI()
diff --git a/src/model.py b/src/model.py
index 7e941dd..977a43a 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=len(BASES), activation="softmax"),
+            Dense(units=32, activation="softmax"),
         ]
     )
     model.compile(