Fixed error in using model.fit output instead of model in probabilities