MichalBusta Update train.py