diff --git a/stargan/attacks.py b/stargan/attacks.py
index fa2acf0..8aa59db 100644
--- a/stargan/attacks.py
+++ b/stargan/attacks.py
@@ -46,7 +46,7 @@ class LinfPGDAttack(object):
 
         self.model.zero_grad()
 
-        return X, eta
+        return X, X - X_nat
 
 def clip_tensor(X, Y, Z):
     # Clip X with Y min and Z max