Add tensor core support to CorrDiff training pipeline

[akshaysubr]

Modulus Pull Request

Description

This PR adds support for TF32 based matrix multiplies to speed up training by ~3.2x

Checklist

• [x] I am familiar with the Contributing Guidelines.
• [ ] New or existing tests cover these changes.
• [x] The documentation is up to date with these changes.
• [ ] The CHANGELOG.md is up to date with these changes.
• [ ] An issue is linked to this pull request.

Dependencies

None