diff --git a/gym_tic_tac_toe/envs/tic_tac_toe_env.py b/gym_tic_tac_toe/envs/tic_tac_toe_env.py
index 8ebaabf..817e30c 100644
--- a/gym_tic_tac_toe/envs/tic_tac_toe_env.py
+++ b/gym_tic_tac_toe/envs/tic_tac_toe_env.py
@@ -33,14 +33,17 @@ class TicTacToeEnv(gym.Env):
 
         # check game over
         for i in range(3):
-            # horizontals and verticals and cross
+            # horizontals and verticals
             if ((board[i * 3] == p and board[i * 3 + 1] == p and board[i * 3 + 2] == p)
-                or (board[i + 0] == p and board[i + 3] == p and board[i + 6] == p)
-                or (board[0] == p and board[4] == p and board[8] == p)
-                or (board[2] == p and board[4] == p and board[6] == p)):
+                or (board[i + 0] == p and board[i + 3] == p and board[i + 6] == p)):
                 reward = p
                 done = True
                 break
+        # diagonals
+        if((board[0] == p and board[4] == p and board[8] == p)
+            or (board[2] == p and board[4] == p and board[6] == p)):
+                reward = p
+                done = True
                 
         return self.state, reward, done, {}
     def _reset(self):