Update SAE/sae.py

SAE/sae.py

@@ -7,6 +7,7 @@ import torch
 import torch.nn as nn
 import os
 import json
+import spaces
 
 class SparseAutoencoder(nn.Module):
     """
@@ -105,6 +106,7 @@ class SparseAutoencoder(nn.Module):
     def n_dirs(self):
         return self.n_dirs_local
 
+    @spaces.GPU
     def encode(self, x):
         x = x - self.pre_bias
         latents_pre_act = self.encoder(x) + self.latent_bias
@@ -120,6 +122,7 @@ 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
@@ -179,7 +182,7 @@ class SparseAutoencoder(nn.Module):
             "auxk_vals": auxk_vals,
         }
 
-    
+    @spaces.GPU
     def decode_sparse(self, inds, vals):
         rows, cols = inds.shape[0], self.n_dirs