Correct the loading snippet for PyTorch >= 2.6

README.md

@@ -45,7 +45,7 @@ attribution built on it inherits that contact.
 | `state_dict` | `OrderedDict` | encoder/decoder weights and `b_pre` |
 | `width` | `int` | 4096 |
 | `k` | `int` | 32 |
-| `mu` | tensor `(1, 512)` | activation mean, subtracted before encoding |
+| `mu` | `np.ndarray (1, 512)` | activation mean, subtracted before encoding |
 | `scale` | `float` | activation scale, divided before encoding |
 | `thresh` | `float` | BatchTopK threshold, frozen over the whole training matrix |
 
@@ -56,16 +56,33 @@ value rather than recomputing per batch.
 
 ## Loading
 
+`mu` is a NumPy array, so on PyTorch 2.6 and later — where `torch.load` defaults
+to `weights_only=True` — the default call raises `UnpicklingError`. Either:
+
 ```python
 import torch
 
-ckpt = torch.load("sae.pt", map_location="cpu")
+ckpt = torch.load("sae.pt", map_location="cpu", weights_only=False)
 print(ckpt["width"], ckpt["k"], ckpt["thresh"])
 
 # Activations are normalised with the stored constants before encoding:
 #   x_norm = (x - ckpt["mu"]) / ckpt["scale"]
 ```
 
+or, to keep the safe loader, allowlist the NumPy reconstructors:
+
+```python
+import numpy as np, torch
+
+with torch.serialization.safe_globals(
+    [np._core.multiarray._reconstruct, np.ndarray, np.dtype, np.dtypes.Float32DType]
+):
+    ckpt = torch.load("sae.pt", map_location="cpu")
+```
+
+On PyTorch below 2.6 the plain `torch.load("sae.pt", map_location="cpu")` works
+as written.
+
 The model class and the encode path are in `src/04_interpretability/train_sae_hcp.py`
 in the code repository below.