diff --git a/genetic_alg.py b/genetic_alg.py
index 59993d9..b542682 100644
--- a/genetic_alg.py
+++ b/genetic_alg.py
@@ -3,9 +3,9 @@ import torch.nn as nn
 import numpy as np
 
 class NeuralNetwork(nn.Module):
-    def __init__(self):
+    def __init__(self, inputSize):
         super(NeuralNetwork, self).__init__()
-        self.hidden = nn.Linear(5, 10)
+        self.hidden = nn.Linear(inputSize, 10)
         self.output = nn.Linear(10, 3)
 
     def forward(self, x):
@@ -14,14 +14,15 @@ class NeuralNetwork(nn.Module):
         return torch.log_softmax(x)
 
 class GeneticAlgorithm:
-    def __init__(self, populationSize: int, mutationRate: float, percentageBest):
+    def __init__(self, populationSize: int, mutationRate: float, percentageBest: float, inputSize: int = 5):
         self.populationSize = populationSize
         self.mutationRate = mutationRate
         self.percentageBest = percentageBest
+        self.inputSize = inputSize
         self.initialize_population()
 
     def initialize_population(self):
-        self.population = [NeuralNetwork() for _ in range(self.populationSize)]
+        self.population = [NeuralNetwork(self.inputSize) for _ in range(self.populationSize)]
     
     def crossover(self, parent1: NeuralNetwork, parent2: NeuralNetwork):
         child1 = NeuralNetwork()
diff --git a/main.py b/main.py
index 9edb7e8..8e71d87 100644
--- a/main.py
+++ b/main.py
@@ -1,6 +1,11 @@
 import pygame
 import sys
 from gameobjects import *
+from genetic_alg import GeneticAlgorithm
+
+POPULATION_SIZE = 50
+MUTATION_RATE = 0.5
+POPULATION_BEST = 0.2
 
 pygame.init()
 
@@ -40,6 +45,8 @@ def get_features(horse):
                HEIGHT, BARRIER_SPEED]
     return features 
 
+genecticAlg = GeneticAlgorithm(POPULATION_SIZE, MUTATION_RATE, POPULATION_BEST)
+
 UI.add_horses(horses)  # Add horses to UI
 while True:
     for event in pygame.event.get():