fix: make torch.load safer with weights_only=True everywhere possible (#624)
Our prepared datasets pickle a lot of our code, and so these are not just
tensors and basic data types, but most of our use of `torch.load()` actually
works properly with `weights_only=True`, so let's use that everywhere possible
and be explicit in the one function where it's not possible.
Partly fixes #621