Fixed model type checking in output probabilities