Fixed bug where app would crash at the end of training