Only learn in 50 epochs for the unit test