add to method

jbloomAus · Feb 28, 2024 · b3f6dc6 · b3f6dc6
1 parent ed8345a
commit b3f6dc6
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/sae_training/sae_group.py b/sae_training/sae_group.py
@@ -45,6 +45,10 @@ def __len__(self):
         # Return the number of SparseAutoencoder instances
         return len(self.autoencoders)
 
+    def to(self, device):
+        for ae in self.autoencoders:
+            ae.to(device)
+
     @classmethod
     def load_from_pretrained(cls, path: str):
         """