feat: update lightning example to lightning 2.0

[dwahdany]

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Docs change / refactoring / dependency upgrade

Motivation and Context / Related issue

Updates the lightning example to use a less hacky way to change datamodules.
Lightning-AI/pytorch-lightning#10430

How Has This Been Tested (if it applies)

Checklist

• The documentation is up-to-date with the changes I made.
• I have read the CONTRIBUTING document and completed the CLA (see CONTRIBUTING).
• All tests passed, and additional code has been covered with new tests.

[facebook-github-bot]

@facebook-github-bot has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[karthikprasad]

Thanks for the PR! There is a minor bug in the variable name. I'm surprised the integration test passed despite this though.

[facebook-github-bot]

@lsc64 has updated the pull request. You must reimport the pull request before landing.

[facebook-github-bot]

@facebook-github-bot has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[dwahdany]

Thanks for the PR! There is a minor bug in the variable name. I'm surprised the integration test passed despite this though.

One could think it’s possible to add two lines without a major oversight, but here we are… thanks, I’ve changed it!

[facebook-github-bot]

@lsc64 has updated the pull request. You must reimport the pull request before landing.

[facebook-github-bot]

@lsc64 has updated the pull request. You must reimport the pull request before landing.

[dwahdany]

friendly ping @karthikprasad
I think your review is implemented. let me know if something else is missing