## Description

Brief description of the changes in this PR.

## 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)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring
- [ ] Test improvements

## Changes Made

- List the main changes made in this PR
- Be specific about what was added, modified, or removed
- Include any new dependencies or configuration changes

## Testing

- [ ] I have added tests that prove my fix is effective or that my feature works
- [ ] New and existing unit tests pass locally with my changes
- [ ] I have run the smoke test (`python -m tests.smoke_test`)
- [ ] I have tested this change on multiple environments (if applicable)

## Documentation

- [ ] I have updated the documentation accordingly
- [ ] I have updated the docstrings for any new or modified functions
- [ ] I have added examples or usage instructions (if applicable)

## Checklist

- [ ] My code follows the style guidelines of this project
- [ ] I have performed a self-review of my own code
- [ ] I have commented my code, particularly in hard-to-understand areas
- [ ] My changes generate no new warnings
- [ ] I have run `pre-commit run --all-files` and fixed any issues
- [ ] I have checked that my changes don't break existing functionality

## Performance Impact

If applicable, describe the performance impact of your changes:
- Memory usage: [increase/decrease/no change]
- Training speed: [faster/slower/no change]
- Inference speed: [faster/slower/no change]

## Breaking Changes

If this is a breaking change, describe:
- What functionality is affected
- How users should migrate their code
- Any deprecation warnings added

## Screenshots/Examples

If applicable, add screenshots or code examples to help explain your changes.

## Related Issues

Closes #[issue number]
Fixes #[issue number]
Related to #[issue number]

## Additional Notes

Any additional information that reviewers should know about this PR.