Bugfix for ot_barycenter with complex numbers

EleutherAI · Sep 19, 2023 · 644ec94 · 644ec94
1 parent bace9ac
commit 644ec94
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/concept_erasure/optimal_transport.py b/concept_erasure/optimal_transport.py
@@ -77,7 +77,7 @@ def ot_barycenter(
         new_loss = mu.trace() + trace_avg - 2 * inner.mul(weights).sum(dim=0).trace()
 
         # Break if the loss is not decreasing
-        if loss - new_loss < tol:
+        if loss.real - new_loss.real < tol:
             break
         else:
             loss = new_loss