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