Provide a better error message on shape mismatches in training loops #480
Description
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
{{ message }}
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