main.py (1 line):
	- line 186: # TODO: move this to train functions