diff --git a/polecart/main/agent.py b/polecart/main/agent.py
index dacc8fd..7fdaf75 100755
--- a/polecart/main/agent.py
+++ b/polecart/main/agent.py
@@ -374,13 +374,13 @@ class Agent:
 					)
 				)
 
+				state = next_state
 
 				# Only train the network if we have enough
 				# transitions in memory to do so.
 				if len(self.memory) >= self.BATCH_SIZE:
 
-					state = next_state
-
+					
 					# Run optimizer
 					self._optimize()