Merge pull request #109 from fabiobrau/main

Bug fix: Projection on the l2 ball
pralab · Nov 2, 2024 · 5eafb33 · 5eafb33
2 parents 1e7f856 + 49653dd
commit 5eafb33
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/secmlt/optimization/constraints.py b/src/secmlt/optimization/constraints.py
@@ -205,7 +205,9 @@ def project(self, x: torch.Tensor) -> torch.Tensor:
         """
         flat_x = x.flatten(start_dim=1)
         diff_norm = flat_x.norm(p=2, dim=1, keepdim=True).clamp_(min=1e-12)
-        flat_x = torch.where(diff_norm <= 1, flat_x, flat_x / diff_norm) * self.radius
+        flat_x = torch.where(
+            diff_norm <= self.radius, flat_x, self.radius * flat_x / diff_norm
+        )
         return flat_x.reshape(x.shape)