Add type hints to improve readability in CustomMLFlowLogger

[princekumarlahon]

Describe your changes

This PR adds type hints to the methods in CustomMLFlowLogger to make the code easier to read and understand.

Type annotations have been added to the __init__, save_dir, and log_image methods based on their expected inputs and outputs.

---

Issue Link

Closes #454

---

Type of change

• 🐛 Bug fix (non-breaking change that fixes an issue)
• ✨ New feature (non-breaking change that adds functionality)
• 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
• 📖 Documentation (Addition or improvements to documentation)

---

Checklist before requesting a review

• My branch is up-to-date with the target branch
• I have performed a self-review of my code
• I have added docstrings for new/modified functions where needed
• I have added comments where necessary
• I have updated the README if required
• I have added tests if needed
• I have given the PR a clear and descriptive title
• I have requested a reviewer/assignee (if applicable)

---

Checklist for reviewers

• the code is readable
• the code is well tested
• the code is documented
• the code is easy to maintain

---

Author checklist after completed review

• I have added a changelog entry if required

---

Checklist for assignee

• PR is up to date with the base branch
• tests pass
• PR is assigned to a milestone (if applicable)
• changelog entry is added

[kshirajahere]

@princekumarlahon Hey this issue is already being solved in #455, please collaborate over there as duplicate PRs consumes extra time of maintainers :)

[joeloskarsson]

See #455

[princekumarlahon]

Thanks for the heads up! I noticed similar changes were already made in earlier PRs. I’ll take this as a learning experience and move on to another issue to contribute. Looking forward to contributing more