It doesn't get very clear that the input dimension of the array doesn't match from this Error. In this case, the input data had an extra dimension trainId.