Add simple FSDP support to MNIST example LightningModule

[amorehead]

What does this PR do?

• Updates the mnist_module.py to reference the Trainer's version of parameters() e.g., for FSDP support
• Note: Without referencing parameters() in this way, Lightning strategies such as FSDP will not be able to successfully wrap one's model parameters. Even more importantly: if one were to train a model while referencing self.parameters() and then attempt to re-train the model when referencing self.trainer.model.parameters(), Lightning 2.0 will (currently) raise an Exception, preventing one from resuming any training with the original checkpoint. That is why I think this change is important for everyone to use as a default.

Before submitting

• Did you make sure title is self-explanatory and the description concisely explains the PR?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you list all the breaking changes introduced by this pull request?
• Did you test your PR locally with pytest command?
• Did you run pre-commit hooks with pre-commit run -a command?

Did you have fun?

Make sure you had fun coding 🙃

[codecov-commenter]

Codecov Report

Patch coverage: 100.00% and no project coverage change.

Comparison is base (1fb5405) 83.24% compared to head (672a3d4) 83.24%.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #604   +/-   ##
=======================================
  Coverage   83.24%   83.24%           
=======================================
  Files          11       11           
  Lines         376      376           
=======================================
  Hits          313      313           
  Misses         63       63           

Files Changed	Coverage Δ
src/models/mnist_module.py	96.96% <100.00%> (ø)

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[ashleve]

Makes sense, thanks for fix!