Add weights_only=True to torch.load() calls

dictionary_learning/activault_s3_buffer.py

@@ -118,7 +118,7 @@ def compile(byte_buffers, shuffle=True, seed=None, return_ids=False):
         # t = torch.from_numpy(n)
         # t = torch.frombuffer(combined_bytes, dtype=dtype) # torch.float32
         buffer = io.BytesIO(combined_bytes)
-        t = torch.load(buffer)
+        t = torch.load(buffer, weights_only=True)
         if (
             isinstance(t, dict) and "states" in t and not return_ids
         ):  # backward compatibility

dictionary_learning/dictionary.py

@@ -129,7 +129,7 @@ def from_pretrained(cls, path, dtype=t.float, device=None, normalize_decoder=Tru
         """
         Load a pretrained autoencoder from a file.
         """
-        state_dict = t.load(path)
+        state_dict = t.load(path, weights_only=True)
         dict_size, activation_dim = state_dict["encoder.weight"].shape
         autoencoder = cls(activation_dim, dict_size)
         autoencoder.load_state_dict(state_dict)
@@ -279,7 +279,7 @@ def from_pretrained(path, device=None):
         """
         Load a pretrained autoencoder from a file.
         """
-        state_dict = t.load(path)
+        state_dict = t.load(path, weights_only=True)
         dict_size, activation_dim = state_dict["encoder.weight"].shape
         autoencoder = GatedAutoEncoder(activation_dim, dict_size)
         autoencoder.load_state_dict(state_dict)
@@ -358,7 +358,7 @@ def from_pretrained(
         loading function.
         """
         if not load_from_sae_lens:
-            state_dict = t.load(path)
+            state_dict = t.load(path, weights_only=True)
             activation_dim, dict_size = state_dict["W_enc"].shape
             autoencoder = JumpReluAutoEncoder(activation_dim, dict_size)
             autoencoder.load_state_dict(state_dict)
@@ -429,7 +429,7 @@ def from_pretrained(path, device=None):
         """
         Load a pretrained autoencoder from a file.
         """
-        state_dict = t.load(path)
+        state_dict = t.load(path, weights_only=True)
         dict_size, activation_dim = state_dict["encoder.weight"].shape
         autoencoder = AutoEncoderNew(activation_dim, dict_size)
         autoencoder.load_state_dict(state_dict)

dictionary_learning/trainers/batch_top_k.py

@@ -79,7 +79,7 @@ def scale_biases(self, scale: float):
 
     @classmethod
     def from_pretrained(cls, path, k=None, device=None, **kwargs) -> "BatchTopKSAE":
-        state_dict = t.load(path)
+        state_dict = t.load(path, weights_only=True)
         dict_size, activation_dim = state_dict["encoder.weight"].shape
         if k is None:
             k = state_dict["k"].item()

dictionary_learning/trainers/matryoshka_batch_top_k.py

@@ -121,7 +121,7 @@ def scale_biases(self, scale: float):
     def from_pretrained(
         cls, path, k=None, device=None, **kwargs
     ) -> "MatryoshkaBatchTopKSAE":
-        state_dict = t.load(path)
+        state_dict = t.load(path, weights_only=True)
         activation_dim, dict_size = state_dict["W_enc"].shape
         if k is None:
             k = state_dict["k"].item()

dictionary_learning/trainers/top_k.py

@@ -137,7 +137,7 @@ def from_pretrained(path, k: Optional[int] = None, device=None):
         """
         Load a pretrained autoencoder from a file.
         """
-        state_dict = t.load(path)
+        state_dict = t.load(path, weights_only=True)
         dict_size, activation_dim = state_dict["encoder.weight"].shape
 
         if k is None: