Update SAE/sae.py

SAE/sae.py

@@ -41,13 +41,6 @@ class SparseAutoencoder(nn.Module):
         self.stats_last_nonzero: torch.Tensor
         self.register_buffer("stats_last_nonzero", torch.zeros(n_dirs_local, dtype=torch.long))
 
-        def auxk_mask_fn(x):
-            dead_mask = self.stats_last_nonzero > dead_steps_threshold
-            x.data *= dead_mask  # inplace to save memory
-            return x
-
-        self.auxk_mask_fn = auxk_mask_fn
-
         ## initialization
 
         # "tied" init
@@ -58,6 +51,11 @@ class SparseAutoencoder(nn.Module):
 
         unit_norm_decoder_(self)
 
+    def auxk_mask_fn(self, x):
+            dead_mask = self.stats_last_nonzero > dead_steps_threshold
+            x.data *= dead_mask  # inplace to save memory
+            return x
+
     def save_to_disk(self, path: str):
         PATH_TO_CFG = 'config.json'
         PATH_TO_WEIGHTS = 'state_dict.pth'
@@ -122,7 +120,6 @@ class SparseAutoencoder(nn.Module):
 
         return latents
 
-    @spaces.GPU
     def forward(self, x):
         x = x - self.pre_bias
         latents_pre_act = self.encoder(x) + self.latent_bias
@@ -182,7 +179,6 @@ class SparseAutoencoder(nn.Module):
             "auxk_vals": auxk_vals,
         }
 
-    @spaces.GPU
     def decode_sparse(self, inds, vals):
         rows, cols = inds.shape[0], self.n_dirs