Fix renorm dimension in random_perturb

MadryLab · Feb 11, 2020 · a733a88 · a733a88
1 parent ca52df7
commit a733a88
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/robustness/attack_steps.py b/robustness/attack_steps.py
@@ -131,7 +131,7 @@ def step(self, x, g):
     def random_perturb(self, x):
         """
         """
-        new_x = x + (ch.rand_like(x) - 0.5).renorm(p=2, dim=1, maxnorm=self.eps)
+        new_x = x + (ch.rand_like(x) - 0.5).renorm(p=2, dim=0, maxnorm=self.eps)
         return ch.clamp(new_x, 0, 1)
 
 # Unconstrained threat model
@@ -152,7 +152,7 @@ def step(self, x, g):
     def random_perturb(self, x):
         """
         """
-        new_x = x + (ch.rand_like(x) - 0.5).renorm(p=2, dim=1, maxnorm=step_size)
+        new_x = x + (ch.rand_like(x) - 0.5).renorm(p=2, dim=0, maxnorm=step_size)
         return ch.clamp(new_x, 0, 1)
 
 class FourierStep(AttackerStep):