lyeoni feat: fix bug in trainer.py