diff --git a/genetic_alg.py b/genetic_alg.py
index b542682..100b14b 100644
--- a/genetic_alg.py
+++ b/genetic_alg.py
@@ -4,14 +4,14 @@ import numpy as np
 
 class NeuralNetwork(nn.Module):
     def __init__(self, inputSize):
-        super(NeuralNetwork, self).__init__()
+        super().__init__()
         self.hidden = nn.Linear(inputSize, 10)
         self.output = nn.Linear(10, 3)
 
     def forward(self, x):
         x = torch.relu(self.hidden(x))
         x = self.output(x)
-        return torch.log_softmax(x)
+        return torch.argmax(x)
 
 class GeneticAlgorithm:
     def __init__(self, populationSize: int, mutationRate: float, percentageBest: float, inputSize: int = 5):
@@ -54,7 +54,11 @@ class GeneticAlgorithm:
             self.population.extend([child1, child2])
         while len(self.population) > self.populationSize: self.population.pop()
     
-    def predict(self, data: list):
+    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]))
diff --git a/main.py b/main.py
index 8e71d87..401797f 100644
--- a/main.py
+++ b/main.py
@@ -45,7 +45,7 @@ def get_features(horse):
                HEIGHT, BARRIER_SPEED]
     return features 
 
-genecticAlg = GeneticAlgorithm(POPULATION_SIZE, MUTATION_RATE, POPULATION_BEST)
+genecticAlg = GeneticAlgorithm(POPULATION_SIZE, MUTATION_RATE, POPULATION_BEST, 10)
 
 UI.add_horses(horses)  # Add horses to UI
 while True:
@@ -56,14 +56,16 @@ while True:
 
     keys = pygame.key.get_pressed()  # Get pressed keys
     
-    for horse in horses:
+    for i, horse in enumerate(horses):
         if horse.stopped == False:  # If the horse is not stopped
-            if keys[pygame.K_UP]:  # Move horse up if UP key is pressed
-                horse.up()           
-            elif keys[pygame.K_DOWN]:  # Move horse down if DOWN key is pressed
-                horse.down()         
+            data = get_features(horse)
+            res = genecticAlg.predict(data, i)
+            if res == 0:
+                horse.up()
+            elif res == 1:
+                horse.down()
             else:
-                horse.stay()  # Stop vertical movement if no key is pressed
+                horse.stay()
         else:
             horse.move(-BARRIER_SPEED, 0)  # Move stopped horse to the left
 
diff --git a/test.py b/test.py
new file mode 100644
index 0000000..6facc2f
--- /dev/null
+++ b/test.py
@@ -0,0 +1,7 @@
+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