diff --git a/genetic_alg.py b/genetic_alg.py
index 100b14b..4492a41 100644
--- a/genetic_alg.py
+++ b/genetic_alg.py
@@ -5,8 +5,8 @@ import numpy as np
 class NeuralNetwork(nn.Module):
     def __init__(self, inputSize):
         super().__init__()
-        self.hidden = nn.Linear(inputSize, 10)
-        self.output = nn.Linear(10, 3)
+        self.hidden = nn.Linear(inputSize, 32)
+        self.output = nn.Linear(32, 3)
 
     def forward(self, x):
         x = torch.relu(self.hidden(x))
@@ -19,16 +19,16 @@ class GeneticAlgorithm:
         self.mutationRate = mutationRate
         self.percentageBest = percentageBest
         self.inputSize = inputSize
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.initialize_population()
 
     def initialize_population(self):
-        self.population = [NeuralNetwork(self.inputSize) for _ in range(self.populationSize)]
+        self.population = [NeuralNetwork(self.inputSize).to(self.device) for _ in range(self.populationSize)]
     
     def crossover(self, parent1: NeuralNetwork, parent2: NeuralNetwork):
-        child1 = NeuralNetwork()
-        child2 = NeuralNetwork()
+        child1 = NeuralNetwork(self.inputSize).to(self.device)
+        child2 = NeuralNetwork(self.inputSize).to(self.device)
         point = len(child1.hidden.weight.data) // 2
-        print([x for x in child1.parameters()])
         child1.hidden.weight.data = torch.cat((parent1.hidden.weight.data[:point], parent2.hidden.weight.data[point:]), dim=0)
         child2.hidden.weight.data = torch.cat((parent2.hidden.weight.data[:point], parent1.hidden.weight.data[point:]), dim=0)
         child1.output.weight.data = parent1.output.weight.data.clone().detach()
@@ -55,11 +55,5 @@ class GeneticAlgorithm:
         while len(self.population) > self.populationSize: self.population.pop()
     
     def predict(self, data: list, i):
-        data = torch.tensor(data, requires_grad=False).float()
-        return self.population[i](data)
-
-    def predict_all(self, data: list):
-        result = []
-        for i in range(self.populationSize):
-            result.append(self.population[i](data[i]))
-        return result
\ No newline at end of file
+        data = torch.tensor(data, requires_grad=False).float().to(self.device)
+        return self.population[i](data)
\ No newline at end of file
diff --git a/main.py b/main.py
index f12dcbf..61e0bb8 100644
--- a/main.py
+++ b/main.py
@@ -3,7 +3,7 @@ import sys
 from gameobjects import *
 from genetic_alg import GeneticAlgorithm
 
-POPULATION_SIZE = 5
+POPULATION_SIZE = 50
 MUTATION_RATE = 0.5
 POPULATION_BEST = 0.2
 FRAME_RATE = 160 #TODO: FIX THE INCORRECT FRAME RATE CORRELATION
@@ -45,6 +45,7 @@ def init_game():
         horse.set_vacceleration(0)
         horse.set_vspeed(0)
         horse.frame_counter = 0
+        horse.fitness = 0
     spawner = Spawner("images/Barrier.png")
     new_barrier = spawner.spawn()
     barriers.append(new_barrier)
@@ -124,6 +125,7 @@ while True:
     
     if all(horse.stopped for horse in horses):
         UI.iteration_num += 1
+        genecticAlg.learn([x.fitness for x in horses])
         init_game()
 
     pygame.display.flip()  # Update the display
diff --git a/test.py b/test.py
deleted file mode 100644
index 6facc2f..0000000
--- a/test.py
+++ /dev/null
@@ -1,7 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as nnf
-import numpy as np
-
-t = torch.tensor([1, 2, 3]).float()
-print(torch.argmax(t))
\ No newline at end of file