Is there some reason why you haven't added the if else part in the last section in the fit-method, ie. like Sentdex has done? 