Fixed ruff formatting issue

EleutherAI · Jan 25, 2024 · 2cff73e · 2cff73e
1 parent 2844e7b
commit 2cff73e
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/concept_erasure/shrinkage.py b/concept_erasure/shrinkage.py
@@ -30,11 +30,11 @@ def optimal_linear_shrinkage(
 
     trace_S = trace(S_n)
 
-    # Since sigma0 is I * tr(S_n) / p, its squared Frobenius norm is just tr(S_n) ** 2 / p.
-    sigma0_norm_sq = trace_S ** 2 / p
+    # Since sigma0 is I * tr(S_n) / p, its squared Frobenius norm is tr(S_n) ** 2 / p.
+    sigma0_norm_sq = trace_S**2 / p
     S_norm_sq = S_n.norm(dim=(-2, -1), keepdim=True) ** 2
 
-    prod_trace = sigma0_norm_sq #torch.linalg.diagonal(S_n) # trace(S_n @ sigma0)
+    prod_trace = sigma0_norm_sq
     top = trace_S * trace_S.conj() * sigma0_norm_sq / n
     bottom = S_norm_sq * sigma0_norm_sq - prod_trace * prod_trace.conj()