fixing

SAE/sae.py

@@ -8,6 +8,7 @@ import torch.nn as nn
 import os
 import json
 import spaces
+import logging
 
 class SparseAutoencoder(nn.Module):
     """
@@ -106,7 +107,6 @@ class SparseAutoencoder(nn.Module):
 
     @spaces.GPU
     def encode(self, x):
-        self = self.to('cuda')
         x = x - self.pre_bias
         latents_pre_act = self.encoder(x) + self.latent_bias
 
@@ -122,6 +122,8 @@ class SparseAutoencoder(nn.Module):
         return latents
 
     def forward(self, x):
+        print(x.device)
+        print(self.pre_bias.device)
         x = x - self.pre_bias
         latents_pre_act = self.encoder(x) + self.latent_bias
         vals, inds = torch.topk(