Add device parameter to fitter.editor

EleutherAI · Jan 10, 2024 · 9b18b3d · 9b18b3d
1 parent b7896b5
commit 9b18b3d
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/concept_erasure/quadratic.py b/concept_erasure/quadratic.py
@@ -146,10 +146,13 @@ def update_single(self, x: Tensor, z: int) -> "QuadraticFitter":
 
  return self
 
- def editor(self) -> QuadraticEditor:
+ def editor(self, device: str | None = None) -> QuadraticEditor:
  """Quadratic editor for the concept."""
  sigma = self.sigma_xx
- return QuadraticEditor(self.mean_x, ot_map(sigma[:, None], sigma))
+ device = device or sigma.device
+
+ T = ot_map(sigma[:, None], sigma).to(device)
+ return QuadraticEditor(self.mean_x.to(device), T)
 
  @cached_property
  def eraser(self) -> QuadraticEraser:

diff --git a/pyproject.toml b/pyproject.toml
@@ -12,7 +12,7 @@ license = {text = "MIT License"}
 dependencies = [
  "torch",
 ]
-version = "0.2.3"
+version = "0.2.4"
 
 [project.optional-dependencies]
 dev = [