fix cosine similarity dimensionality check (#66191)
Summary:
Fixes https://github.com/pytorch/pytorch/issues/66086
Pull Request resolved: https://github.com/pytorch/pytorch/pull/66191
Reviewed By: dagitses, malfet
Differential Revision: D31436997
Pulled By: ngimel
fbshipit-source-id: 363556eea4e1696d928ae08320d298451c286b10
