Add checkpoint tracking for loading

[one-some]

Torch only. Not sure if we even have support for loading multi-checkpoint safetensors anyway, since its sort of auto-sharded per tensor