Improve torch.cdist performance (#20605)
Summary:
Fix based on https://github.com/pytorch/pytorch/issues/15253
Pull Request resolved: https://github.com/pytorch/pytorch/pull/20605
Differential Revision: D15396123
Pulled By: ifedan
fbshipit-source-id: 3ed373e68339a35360f083d4aad1b655abcaf97e
