diff --git a/src/cortecs/tests/test_integration.py b/src/cortecs/tests/test_integration.py
index 47e853a..eb6a48d 100644
--- a/src/cortecs/tests/test_integration.py
+++ b/src/cortecs/tests/test_integration.py
@@ -202,7 +202,7 @@ def test_optimize(self):
             optimizer.best_params
             == {
                 "n_layers": 2,
-                "n_neurons": 2.0,
+                "n_neurons": 2,
                 "activation": "sigmoid",
                 "learn_rate": 0.01,
             }