I found the issue in the "min_prediction_length" parameter. Initially I set it same as the max_prediction_length , but after changing it to a lower value, model was working fine.