paranthesis tricks bjork

deel-ai · Jul 1, 2024 · 0cbbf2d · 0cbbf2d
1 parent 4583add
commit 0cbbf2d
Showing 1 changed file with 11 additions and 4 deletions.
diff --git a/deel/torchlip/normalizers.py b/deel/torchlip/normalizers.py
@@ -70,10 +70,17 @@ def bjorck_normalization(
     shape = w.shape
     cout = w.size(0)
     w_mat = w.reshape(cout, -1)
-    for i in range(niter):
-        w_mat = (1.0 + beta) * w_mat - beta * torch.mm(
-            w_mat, torch.mm(w_mat.t(), w_mat)
-        )
+
+    if w_mat.shape[0]>w_mat.shape[1]:
+        for i in range(niter):
+            w_mat = (1.0 + beta) * w_mat - beta * torch.mm(
+                w_mat, torch.mm(w_mat.t(), w_mat)
+            )
+    else:
+        for i in range(niter):
+            w_mat = (1.0 + beta) * w_mat - beta * torch.mm(
+                torch.mm(w_mat, w_mat.t()), w_mat
+            )
     w = w_mat.reshape(shape)
     return w