Fix disk offload for full safetensors checkpoints
# What does this PR do?

#20321 was only tested with safetensors checkpoints containing multiple shards. The code failed for full checkpoints, this PR fixes it.
_The documentation is not available anymore as the PR was closed or merged._