refactor train.py

[masoud-khalilian]

Make the code in train.py more readable and modular by adding a function with an explicit name; reading the code would be easier and remove unnecessary clutter from the code.