Configured for computing cluster

ai-safety-foundation · Apr 1, 2024 · 2660f57 · 2660f57
1 parent cf79a57
commit 2660f57
Show file tree

Hide file tree

Showing 5 changed files with 144 additions and 57 deletions.
diff --git a/docs/content/demo.ipynb b/docs/content/demo.ipynb
@@ -424,7 +424,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": ".venv",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -438,7 +438,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.4"
+   "version": "3.10.9"
   },
   "vscode": {
    "interpreter": {
@@ -447,5 +447,5 @@
   }
  },
  "nbformat": 4,
- "nbformat_minor": 2
+ "nbformat_minor": 4
 }
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -17,8 +17,8 @@
         python=">=3.10, <3.12"
         strenum=">=0.4.15"
         tokenizers=">=0.15.1"
-        torch=">=2.1.2"
-        torchmetrics="^1.3.0.post0"
+        torch="2.1.2"
+        torchmetrics="1.3.0.post0"
         transformers=">=4.37.2"
         wandb=">=0.16.2"
         zstandard=">=0.22.0"

diff --git a/sparse_autoencoder/train/sweep.py b/sparse_autoencoder/train/sweep.py
@@ -76,7 +76,7 @@ def setup_autoencoder(
         type=type,
     )
 
-    return LitSparseAutoencoder(config).cuda()
+    return LitSparseAutoencoder(config)
 
 
 def setup_source_data(hyperparameters: RuntimeHyperparameters) -> SourceDataset: