Issue 1:
🐛 Bug

In PyTorch 1.5 (not 1.4) using cdist produces NANs instead of regular gradients. This happens on both CPU and GPU.

To Reproduce

Steps to reproduce the behavior:

Download the tensors

Run the code

import torch

emb1, emb2, cdist_grad = torch.load('cdist_grad.pt')

emb1.retain_grad()
d = torch.cdist(emb1, emb2)
d.backward(cdist_grad)

print(emb1.grad[0, 17])

The printed gradients will be NAN.

Expected behavior

Actual numerical gradients.

Environment

PyTorch Version (e.g., 1.0): 1.5.0

OS (e.g., Linux): Linux

How you installed PyTorch (conda, pip, source): conda

Build command you used (if compiling from source):

Python version: 3.7.6

CUDA/cuDNN version: 10.2

GPU models and configuration: Nvidia 1080 Ti (but it also happens on CPU)

Any other relevant information:

Issue 2:
🐛 Bug

torch.cdist produces nan gradients in Pytorch 1.5, but not Pytorch 1.4

To Reproduce

Steps to reproduce the behavior:

Run the similarity example here: https://github.com/pytorch/vision/tree/master/references/similarity

On Pytorch 1.4, training proceeds normally with the loss decreasing each iteration

On Pytorch 1.5, the gradients become nan after the first iteration, which breaks training

I used torch.autograd.detect_anomaly to identify that torch.cdist is the problem (error trace below)

---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
~/Projects/vision/references/similarity/train.py in <module>
171 if __name__ == '__main__':
172 args = parse_args()
--> 173 main(args)

~/Projects/vision/references/similarity/train.py in main(args)
132 for epoch in range(1, args.epochs + 1):
133 print('Training...')
--> 134 train_epoch(model, optimizer, criterion, train_loader,
135 device, epoch, args.print_freq)
136

~/Projects/vision/references/similarity/train.py in train_epoch(model, optimizer, criterion, data_loader, device, epoch, print_freq)
25 with torch.autograd.detect_anomaly():
26 loss, frac_pos_triplets = criterion(embeddings, targets)
---> 27 loss.backward()
28 optimizer.step()
29

~/miniconda3/envs/pytorch/lib/python3.8/site-packages/torch/tensor.py in backward(self, gradient, retain_graph, create_graph)
196 products. Defaults to ``False``.
197 """
--> 198 torch.autograd.backward(self, gradient, retain_graph, create_graph)
199
200 def register_hook(self, hook):

~/miniconda3/envs/pytorch/lib/python3.8/site-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables)
96 retain_graph = create_graph
97
---> 98 Variable._execution_engine.run_backward(
99 tensors, grad_tensors, retain_graph, create_graph,
100 allow_unreachable=True) # allow_unreachable flag

RuntimeError: Function 'SqrtBackward' returned nan values in its 0th output.

Expected behavior

Training proceeds normally without nan gradients due to torch.cdist

Environment

PyTorch Version (e.g., 1.0): Pytorch 1.5

OS (e.g., Linux): Ubuntu 18.04 LTS

How you installed PyTorch (conda, pip, source): conda install pytorch torchvision cudatoolkit=10.1 -c pytorch

Python version: 3.8.2

CUDA/cuDNN version: 10.1/7.6

GPU models and configuration: RTX 2080 Ti/Titan Xp (tried both)

Any other relevant information: I believe things work in Pytorch 1.4

cc @ezyang @gchanan @zou3519 @fmassa