fix: allow weights_only=False when loading --cache-dataset torch cache

[musicalplatypus]

Problem

load_data() in tinyml-tinyverse/tinyml_tinyverse/common/utils/utils.py has two cache-hit branches (training dataset and validation/test dataset) that call torch.load(cache_path) with no weights_only argument. The cached object is a (dataset, datadir) tuple where dataset is one of this project's own Dataset subclasses.

PyTorch 2.6+ changed torch.load()'s default to weights_only=True, which only allowlists a small set of known-safe types. Any run that reuses a --cache-dataset cache written by an earlier run now crashes with:

_pickle.UnpicklingError: Weights only load failed...
WeightsUnpickler error: Unsupported global: GLOBAL <DatasetClass>
was not an allowed global by default...

i.e. the --cache-dataset speedup flag is currently broken on PyTorch 2.6+.

Fix

Pass weights_only=False explicitly at both call sites. This is not a security fix — cache_path is derived from a hash of datadir and is a purely local cache this same process wrote moments earlier; it is never fetched from a URL or otherwise externally supplied. This is a different situation from genuinely untrusted/externally-sourced checkpoint loads elsewhere in the codebase, which should keep the safe default or use torch.serialization.safe_globals([...]).

Testing

Added tinyml-tinyverse/tests/test_cache_dataset_weights_only.py, which reproduces the crash end-to-end through load_data() itself (writes a fake cached dataset via torch.save, then calls load_data() against it), not just a direct torch.load call.

• Verified the test fails on pre-fix code with the exact UnpicklingError above (checked via git stash).
• Verified the test passes with the fix.

Merge safety

This touches the same file (utils.py) as three other currently-open PRs (pr/mps-eval-fixes, pr/compile-hardening, pr/hardware-defaults). Verified via a disposable worktree that this branch merges cleanly on top of all three with no conflicts (git merge --no-ff of all three followed by this branch resolves via auto-merge, clean working tree).

🤖 Generated with Claude Code

[musicalplatypus]

An independent peer review of this PR's weights_only fix caught a deeper, pre-existing bug it had unblocked (previously masked by the crash). Pushed a follow-up commit (de8ec33):

_get_cache_path(datadir) computed the same cache path for the training and validation datasets (both derived only from datadir), and only the training branch ever wrote its cache -- the validation write-back was dead code, commented out with # TODO: Add utils and uncomment the if block.

Consequence, confirmed by reproduction: on the very first run with --cache-dataset enabled, load_data() writes the training dataset to cache_path, then re-derives the identical path for validation and finds it already exists -- so dataset_test silently becomes a torch.load() of the training dataset. Not a crash: --cache-dataset looked like it worked, but every validation metric from that point on was computed against training data, not held-out data.

Fixed by giving _get_cache_path a tag parameter ('train'/'val') folded into the hash, and by enabling the validation cache write-back that was dead code.

Also rewrote the test file: the previous test patched _get_cache_path with a single fixed return value for both train and val calls, which (ironically) would have passed under this exact bug. New tests use a tag-aware mock and assert train/val come back distinct -- verified to fail on the pre-fix code with the exact symptom and pass post-fix. Re-verified this still merges cleanly alongside PRs #21/#22/#23 (all touch utils.py).