Fix selective restore by explicitly verifying keys
<!-- pylint-output-strict-mode-false -->

beep boop 🤖: 🙏 The following files have warnings. In case you are familiar with these, please try helping us to improve the code base.

---

Your code was analyzed with PyLint. The following annotations have been identified:

```

------------------------------------
Your code has been rated at 10.00/10
```

---

Thank you for improving NeMo's documentation! [🤖]: Hi @hemildesai 👋,

We wanted to let you know that a [CICD pipeline](https://github.com/NVIDIA/NeMo/actions/runs/11977485168) for this PR just finished successfully

So it might be time to merge this PR or get some approvals

I'm just a bot so I'll leave it you what to do next.

//cc @pablo-garay @ko3n1g
io.load_checkpoint always loads the entire state including `optimizer`. For selective restore, we need to explicitly select keys  to restore.