fixed dimensions issue with the iris dataset

neural_network/backprop.py

@@ -48,8 +48,9 @@ def bp(
             error * model.func_prime(y_hat),
         )
         db2 = np.sum(error * model.func_prime(y_hat), axis=0)
-        db1 = np.sum(np.dot(error * model.func_prime(y_hat), model.w2.T)
-                     * model.func_prime(node1), axis=0)
+        db1 = np.sum(
+            np.dot(error * model.func_prime(y_hat), model.w2.T) * model.func_prime(node1), axis=0,
+        )
 
         # update weights & biases using gradient descent.
         # this is -= and not += because if the gradient descent

neural_network/main.py

@@ -5,10 +5,7 @@ from neural_network.opts import activation
 from neural_network.backprop import bp
 
 
-def init(
-    X: np.array,
-    hidden_size: int
-) -> dict:
+def init(X: np.array, hidden_size: int) -> dict:
     """
     returns a dictionary containing randomly initialized
     weights and biases to start off the neural_network
@@ -16,8 +13,8 @@ def init(
     return {
         "w1": np.random.randn(X.shape[1], hidden_size),
         "b1": np.zeros((1, hidden_size)),
-        "w2": np.random.randn(hidden_size, 1),
-        "b2": np.zeros((1, 1)),
+        "w2": np.random.randn(hidden_size, 3),  # Output layer has 3 neurons
+        "b2": np.zeros((1, 3)),  # Output layer has 3 neurons
     }