mypy: make optimization module pass type checks (#1729)

[nadae1010]

Summary

This PR fixes type errors in the optimization module as part of the mypy-compliance effort (#1719). Addresses #1729.

Changes

• Widened OptimizerConfig.build signature to accept *args, **kwargs.
• Allowed Optional[int] for num_warmup_steps in LRSchedulerConfig.
• Verified with mypy: optimization module now passes cleanly.

Before vs After (local)

• Before: 5 errors (see /tmp/mypy-before.txt)
• After: 0 errors (see /tmp/mypy-after.txt)

[AdilZouitine]

Hey could you uncomment the mypy hook in .pre-commit-config.yaml and uncomment the module in pyproject.toml apart that it's LGTM 😄

[imstevenpmwork]

Hello @nadae1010,

Can you tag me whenever this is ready for a review ? Feel free to pin me in discord too
Sorry for the back and forth, thanks !