Improve dimensionality mismatch warning (#47874)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/47874
Test Plan: N/A
Reviewed By: ngimel
Differential Revision: D24926123
fbshipit-source-id: ace5543ae5122906164e13ae9463fe4dfa74d8d6
