rename var

stanfordnlp · Jan 21, 2024 · 78b3ac5 · 78b3ac5
1 parent 48d1a3b
commit 78b3ac5
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/pyvene/models/interventions.py b/pyvene/models/interventions.py
@@ -532,11 +532,12 @@ def __init__(self, embed_dim, **kwargs):
         self.interchange_dim = embed_dim
         rs = np.random.RandomState(1)
         prng = lambda *shape: rs.randn(*shape)
-        self.noise = torch.from_numpy(
-            prng(1, 4, embed_dim)).to(device)
-        self.noise_level = kwargs["noise_leve"] \
+        noise_level = kwargs["noise_leve"] \
             if "noise_leve" in kwargs else 0.13462981581687927 
-
+        self.register_buffer('noise', torch.from_numpy(
+            prng(1, 4, embed_dim)))
+        self.register_buffer('noise_level', torch.tensor(noise_level))
+
     def forward(self, base, source=None, subspaces=None):
         base[..., : self.interchange_dim] += self.noise * self.noise_level
         return base

diff --git a/pyvene_101.ipynb b/pyvene_101.ipynb
@@ -1583,7 +1583,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 100,
+   "execution_count": 2,
    "id": "c0b6a70f",
    "metadata": {},
    "outputs": [
@@ -1606,7 +1606,7 @@
     "    [{\"component\": \"block_input\"}] + \n",
     "    [{\"layer\": l, \"component\": c} \n",
     "      for l in range(s, e)],\n",
-    "    [pv.ZeroIntervention] +\n",
+    "    [pv.NoiseIntervention] +\n",
     "    [pv.VanillaIntervention]*(e-s))\n",
     "  return config\n",
     "\n",