Merge pull request #7 from EleutherAI/where-fix

Use torch.where instead of Tensor.where for PyTorch <2.0
EleutherAI · Sep 7, 2023 · 9689ef5 · 9689ef5
2 parents 0d4dbcb + 5866ce4
commit 9689ef5
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/concept_erasure/leace.py b/concept_erasure/leace.py
@@ -192,8 +192,8 @@ def eraser(self) -> LeaceEraser:
  # Assuming PSD; account for numerical error
  L.clamp_min_(0.0)
 
- W = V * L.rsqrt().where(mask, 0.0) @ V.mT
- W_inv = V * L.sqrt().where(mask, 0.0) @ V.mT
+ W = V * torch.where(mask, L.rsqrt(), 0.0) @ V.mT
+ W_inv = V * torch.where(mask, L.sqrt(), 0.0) @ V.mT
  else:
  W, W_inv = eye, eye
 

diff --git a/concept_erasure/optimal_transport.py b/concept_erasure/optimal_transport.py
@@ -28,7 +28,7 @@ def psd_sqrt_rsqrt(A: Tensor) -> tuple[Tensor, Tensor]:
  # We actually compute the pseudo-inverse here for numerical stability.
  # Use the same heuristic as `torch.linalg.pinv` to determine the tolerance.
  thresh = L[..., None, -1] * A.shape[-1] * torch.finfo(A.dtype).eps
- rsqrt = U * L.rsqrt().where(L > thresh, 0.0) @ U.mT
+ rsqrt = U * torch.where(L > thresh, L.rsqrt(), 0.0) @ U.mT
 
  return sqrt, rsqrt