Fix parallel_for when grain_size > 0

[Achazwl]

This do not affect current BMTrain since BMTrain pass grain_size=0 in all parallel_for.

Pull Request Template

Issue Reference

Please mention the issue number if applicable, or write "N/A" if it's a new feature.

Issue #...

Description

Please describe your changes in detail. If it resolves an issue, please state how it resolves it.

Type of Change

• [ ] 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 not work as expected)
• [ ] This change requires a documentation update

How Has This Been Tested?

Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce.

Checklist

• [ ] I have read the CONTRIBUTING document.
• [ ] My code follows the code style of this project.
• [ ] My change requires a change to the documentation.
• [ ] I have updated the documentation accordingly.
• [ ] I have added tests to cover my changes.
• [ ] All new and existing tests passed.

Additional Information

Any additional information, configuration, or data that might be necessary for the review.