fixing train test split issue with iris dataset

dataset/iris.py

@@ -1,6 +1,5 @@
 from sklearn.datasets import load_iris
 from sklearn.preprocessing import OneHotEncoder, StandardScaler
-from sklearn.model_selection import train_test_split
 import numpy as np
 
 
@@ -11,19 +10,7 @@ def iris() -> tuple[np.array]:
     after being normalized and one-hot encoded 
     """
     iris = load_iris()
-    X_train, X_test, y_train, y_test = train_test_split(
-        iris.data,
-        iris.target,
-        test_size=0.3,
-        random_state=8675309,
-    )
     scaler = StandardScaler()
-    X_train, X_test = scaler.fit_transform(
-        X_train
-    ), scaler.fit_transform(
-        X_test
-    )
-
-    y_train = OneHotEncoder().fit_transform(y_train.reshape(-1, 1)).toarray()
-    y_test = OneHotEncoder().fit_transform(y_test.reshape(-1, 1)).toarray()
-    return X_train, X_test, y_train, y_test
+    x = scaler.fit_transform(iris.data)
+    y = OneHotEncoder().fit_transform(iris.target.reshape(-1, 1)).toarray()
+    return x, y

neural_network/main.py

@@ -22,16 +22,21 @@ def init(
 
 
 def main( 
-    X_train: np.array,
-    y_train: np.array,
-    X_test: np.array,
-    y_test: np.array,
+    X: np.array,
+    y: np.array,
     args,
 ) -> None:
-    wb = init(X_train, args["hidden_size"])
+    wb = init(X, args["hidden_size"])
     act = activation[args["activation_func"]]
     args["activation_func"] = act["main"]
     args["func_prime"] = act["prime"]
+    X_train, X_test, y_train, y_test = train_test_split(
+        X,
+        y,
+        test_size=0.2,
+        random_state=8675309,
+    )
+
     model = bp(X_train, y_train, wb, args)
 
     # evaluate the model and return final results