Make loading device agnostic

ai-safety-foundation · Jan 26, 2024 · 39a3d3a · 39a3d3a
1 parent f55652e
commit 39a3d3a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sparse_autoencoder/autoencoder/model.py b/sparse_autoencoder/autoencoder/model.py
@@ -276,7 +276,7 @@ def load(
             The loaded model.
         """
         # Load the file
-        serialized_state = torch.load(file_path, map_location=torch.device("cpu"))
+        serialized_state = torch.load(file_path)
         state = SparseAutoencoderState.model_validate(serialized_state)
 
         # Initialise the model