Pull changes that Hamze made for his PR #2380.

[DongchenZ]

Description

Motivation and Context

Review Time Estimate

• [ ] Immediately
• [ ] Within one week
• [ ] When possible

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)

Checklist:

• [ ] My change requires a change to the documentation.
• [ ] My name is in the list of CITATION.cff
• [ ] I have updated the CHANGELOG.md.
• [ ] I have updated the documentation accordingly.
• [ ] I have read the CONTRIBUTING document.
• [ ] I have added tests to cover my changes.
• [ ] All new and existing tests passed.