perf(aggregation): Prevent cuda sync in normalize

Author: ValerianRey

src/torchjd/_linalg/_gramian.py

@@ -52,10 +52,8 @@ def normalize(gramian: PSDMatrix, eps: float) -> PSDMatrix:
     therefore `G` divided by the sum of its diagonal elements.
     """
     squared_frobenius_norm = gramian.diagonal().sum()
-    if squared_frobenius_norm < eps:
-        output = torch.zeros_like(gramian)
-    else:
-        output = gramian / squared_frobenius_norm
+    condition = squared_frobenius_norm < eps
+    output = torch.where(condition, torch.zeros_like(gramian), gramian / squared_frobenius_norm)
     return cast(PSDMatrix, output)