diff --git a/src/llmart/attack.py b/src/llmart/attack.py
index e85a629..d34d58e 100644
--- a/src/llmart/attack.py
+++ b/src/llmart/attack.py
@@ -290,8 +290,8 @@ def train(
             pbar.set_postfix(postfix)
 
             # Save tokens with highest success rate
-            if success_rate > best_success_rate:
-                best_success_rate, best_attack = loss, copy.deepcopy(attack)
+            if success_rate >= best_success_rate:
+                best_success_rate, best_attack = success_rate, copy.deepcopy(attack)
 
             # Exit attack loop if we found a successful attack across all training examples
             if (