Provide a better error message on shape mismatches in training loops

[seanmor5]

Right now we just raise the Nx error, but we should do a better job of making it clear why the mismatch error occurs and possible options to fix it