fix(trainer): custom trainer

[LongxingTan]

Description

This PR ...

Checklist

• Linked issues (if existing)
• Amended changelog.md for large changes (and added myself there as contributor)
• Added/modified tests
• Used pre-commit hooks when committing to ensure that code is compliant with hooks. Install hooks with pre-commit install.
  To run hooks independent of commit, execute pre-commit run --all-files

Thank you for joining. Have fun coding!

[codecov]

Codecov Report

Attention: Patch coverage is 84.49612% with 20 lines in your changes missing coverage. Please review.

Project coverage is 55.37%. Comparing base (bcbf148) to head (30a69cb).
Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
src/retrievals/trainer/custom_trainer.py	83.47%	20 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master      #87      +/-   ##
==========================================
+ Coverage   51.98%   55.37%   +3.39%     
==========================================
  Files          38       38              
  Lines        2826     2725     -101     
==========================================
+ Hits         1469     1509      +40     
+ Misses       1357     1216     -141     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.