fix mps bug

jbloomAus · Mar 22, 2024 · e7b238f · e7b238f
1 parent 4c03b3d
commit e7b238f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sae_training/sae_group.py b/sae_training/sae_group.py
@@ -69,7 +69,7 @@ def load_from_pretrained(cls, path: str):
             try:
                 if torch.backends.mps.is_available():
                     group = torch.load(path, map_location="mps")
-                    group.cfg.device = "mps"
+                    group["cfg"].device = "mps"
                 else:
                     group = torch.load(path)
             except Exception as e: