Fix Distributed Checkpoint Tutorial

[LTMeyer]

Fixes #3067

Description

Update the code provided in the tutorial.

• Fix undefined variables model and optimizer;
• Fix missing import of Stateful;
• Fix the state_dict variable for correct loading.

Checklist

• The issue that is being fixed is referred in the description (see above "Fixes #ISSUE_NUMBER")
• Only one issue is addressed in this pull request
• Labels from the issue that this PR is fixing are added to this pull request
• No unnecessary issues are included into this pull request.

cc @wconstab @osalpekar @H-Huang @kwen2501

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/tutorials/3068

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• PyTorch Testing Nodes Undergoing ROCm 6.2.1 Upgrades

✅ No Failures

As of commit 1f1aad6 with merge base be7f1b3 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.