<h2 dir="auto">🐛 Bug</h2>
<p dir="auto">In PyTorch 1.5 (not 1.4) using cdist produces NANs instead of regular gradients. This happens on both CPU and GPU.</p>
<h2 dir="auto">To Reproduce</h2>
<p dir="auto">Steps to reproduce the behavior:</p>
<ol dir="auto">
<li><a href="https://figshare.com/s/47d6035b990fa62d730d" rel="nofollow">Download the tensors</a></li>
<li>Run the code</li>
</ol>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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])"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">torch</span>

<span class="pl-s1">emb1</span>, <span class="pl-s1">emb2</span>, <span class="pl-s1">cdist_grad</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">load</span>(<span class="pl-s">'cdist_grad.pt'</span>)

<span class="pl-s1">emb1</span>.<span class="pl-en">retain_grad</span>()
<span class="pl-s1">d</span> <span class="pl-c1">=</span> <span class="pl-s1">torch</span>.<span class="pl-en">cdist</span>(<span class="pl-s1">emb1</span>, <span class="pl-s1">emb2</span>)
<span class="pl-s1">d</span>.<span class="pl-en">backward</span>(<span class="pl-s1">cdist_grad</span>)

<span class="pl-en">print</span>(<span class="pl-s1">emb1</span>.<span class="pl-s1">grad</span>[<span class="pl-c1">0</span>, <span class="pl-c1">17</span>])</pre></div>
<ol start="3" dir="auto">
<li>The printed gradients will be NAN.</li>
</ol>

<h2 dir="auto">Expected behavior</h2>

<p dir="auto">Actual numerical gradients.</p>
<h2 dir="auto">Environment</h2>
<ul dir="auto">
<li>PyTorch Version (e.g., 1.0): 1.5.0</li>
<li>OS (e.g., Linux): Linux</li>
<li>How you installed PyTorch (<code class="notranslate">conda</code>, <code class="notranslate">pip</code>, source): conda</li>
<li>Build command you used (if compiling from source):</li>
<li>Python version: 3.7.6</li>
<li>CUDA/cuDNN version: 10.2</li>
<li>GPU models and configuration: Nvidia 1080 Ti (but it also happens on CPU)</li>
<li>Any other relevant information:</li>
</ul>
<h2 dir="auto"><g-emoji class="g-emoji" alias="bug" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/1f41b.png">🐛</g-emoji> Bug</h2>
<p dir="auto">torch.cdist produces nan gradients in Pytorch 1.5, but not Pytorch 1.4</p>
<h2 dir="auto">To Reproduce</h2>
<p dir="auto">Steps to reproduce the behavior:</p>
<ol dir="auto">
<li>
<p dir="auto">Run the similarity example here: <a href="https://github.com/pytorch/vision/tree/master/references/similarity">https://github.com/pytorch/vision/tree/master/references/similarity</a></p>
</li>
<li>
<p dir="auto">On Pytorch 1.4, training proceeds normally with the loss decreasing each iteration</p>
</li>
<li>
<p dir="auto">On Pytorch 1.5, the gradients become nan after the first iteration, which breaks training</p>
</li>
<li>
<p dir="auto">I used torch.autograd.detect_anomaly to identify that torch.cdist is the problem (error trace below)</p>
</li>
</ol>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
~/Projects/vision/references/similarity/train.py in &lt;module&gt;
    171 if __name__ == '__main__':
    172     args = parse_args()
--&gt; 173     main(args)

~/Projects/vision/references/similarity/train.py in main(args)
    132     for epoch in range(1, args.epochs + 1):
    133         print('Training...')
--&gt; 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)
---&gt; 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         &quot;&quot;&quot;
--&gt; 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 
---&gt; 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."><pre class="notranslate"><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span>
<span class="pl-v">RuntimeError</span>                              <span class="pl-v">Traceback</span> (<span class="pl-s1">most</span> <span class="pl-s1">recent</span> <span class="pl-s1">call</span> <span class="pl-s1">last</span>)
<span class="pl-c1">~</span><span class="pl-c1">/</span><span class="pl-v">Projects</span><span class="pl-c1">/</span><span class="pl-s1">vision</span><span class="pl-c1">/</span><span class="pl-s1">references</span><span class="pl-c1">/</span><span class="pl-s1">similarity</span><span class="pl-c1">/</span><span class="pl-s1">train</span>.<span class="pl-s1">py</span> <span class="pl-c1">in</span> <span class="pl-c1">&lt;</span><span class="pl-s1">module</span><span class="pl-c1">&gt;</span>
    <span class="pl-c1">171</span> <span class="pl-k">if</span> <span class="pl-s1">__name__</span> <span class="pl-c1">==</span> <span class="pl-s">'__main__'</span>:
    <span class="pl-c1">172</span>     <span class="pl-s1">args</span> <span class="pl-c1">=</span> <span class="pl-en">parse_args</span>()
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">173</span>     <span class="pl-en">main</span>(<span class="pl-s1">args</span>)

<span class="pl-c1">~</span><span class="pl-c1">/</span><span class="pl-v">Projects</span><span class="pl-c1">/</span><span class="pl-s1">vision</span><span class="pl-c1">/</span><span class="pl-s1">references</span><span class="pl-c1">/</span><span class="pl-s1">similarity</span><span class="pl-c1">/</span><span class="pl-s1">train</span>.<span class="pl-s1">py</span> <span class="pl-c1">in</span> <span class="pl-en">main</span>(<span class="pl-s1">args</span>)
    <span class="pl-c1">132</span>     <span class="pl-s1">for</span> <span class="pl-s1">epoch</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-c1">1</span>, <span class="pl-s1">args</span>.<span class="pl-s1">epochs</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>):
    <span class="pl-c1">133</span>         <span class="pl-en">print</span>(<span class="pl-s">'Training...'</span>)
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">134</span>         <span class="pl-en">train_epoch</span>(<span class="pl-s1">model</span>, <span class="pl-s1">optimizer</span>, <span class="pl-s1">criterion</span>, <span class="pl-s1">train_loader</span>,
    <span class="pl-c1">135</span>                     <span class="pl-s1">device</span>, <span class="pl-s1">epoch</span>, <span class="pl-s1">args</span>.<span class="pl-s1">print_freq</span>)
    <span class="pl-c1">136</span> 

<span class="pl-c1">~</span><span class="pl-c1">/</span><span class="pl-v">Projects</span><span class="pl-c1">/</span><span class="pl-s1">vision</span><span class="pl-c1">/</span><span class="pl-s1">references</span><span class="pl-c1">/</span><span class="pl-s1">similarity</span><span class="pl-c1">/</span><span class="pl-s1">train</span>.<span class="pl-s1">py</span> <span class="pl-c1">in</span> <span class="pl-en">train_epoch</span>(<span class="pl-s1">model</span>, <span class="pl-s1">optimizer</span>, <span class="pl-s1">criterion</span>, <span class="pl-s1">data_loader</span>, <span class="pl-s1">device</span>, <span class="pl-s1">epoch</span>, <span class="pl-s1">print_freq</span>)
     <span class="pl-c1">25</span>         <span class="pl-k">with</span> <span class="pl-s1">torch</span>.<span class="pl-s1">autograd</span>.<span class="pl-en">detect_anomaly</span>():
     <span class="pl-c1">26</span>             <span class="pl-s1">loss</span>, <span class="pl-s1">frac_pos_triplets</span> <span class="pl-c1">=</span> <span class="pl-en">criterion</span>(<span class="pl-s1">embeddings</span>, <span class="pl-s1">targets</span>)
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">27</span>             <span class="pl-s1">loss</span>.<span class="pl-en">backward</span>()
     <span class="pl-c1">28</span>         <span class="pl-s1">optimizer</span>.<span class="pl-en">step</span>()
     <span class="pl-c1">29</span> 

<span class="pl-c1">~</span><span class="pl-c1">/</span><span class="pl-s1">miniconda3</span><span class="pl-c1">/</span><span class="pl-s1">envs</span><span class="pl-c1">/</span><span class="pl-s1">pytorch</span><span class="pl-c1">/</span><span class="pl-s1">lib</span><span class="pl-c1">/</span><span class="pl-s1">python3</span>.<span class="pl-c1">8</span><span class="pl-c1">/</span><span class="pl-s1">site</span><span class="pl-c1">-</span><span class="pl-s1">packages</span><span class="pl-c1">/</span><span class="pl-s1">torch</span><span class="pl-c1">/</span><span class="pl-s1">tensor</span>.<span class="pl-s1">py</span> <span class="pl-c1">in</span> <span class="pl-en">backward</span>(<span class="pl-s1">self</span>, <span class="pl-s1">gradient</span>, <span class="pl-s1">retain_graph</span>, <span class="pl-s1">create_graph</span>)
    <span class="pl-c1">196</span>                 <span class="pl-s1">products</span>. <span class="pl-v">Defaults</span> <span class="pl-s1">to</span> <span class="pl-s">``</span><span class="pl-v">False</span><span class="pl-s">``</span>.
    <span class="pl-c1">197</span>         """
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">198</span>         <span class="pl-s1">torch</span>.<span class="pl-s1">autograd</span>.<span class="pl-en">backward</span>(<span class="pl-s1">self</span>, <span class="pl-s1">gradient</span>, <span class="pl-s1">retain_graph</span>, <span class="pl-s1">create_graph</span>)
    <span class="pl-c1">199</span> 
    <span class="pl-c1">200</span>     <span class="pl-k">def</span> <span class="pl-en">register_hook</span>(<span class="pl-s1">self</span>, <span class="pl-s1">hook</span>):

<span class="pl-c1">~</span><span class="pl-c1">/</span><span class="pl-s1">miniconda3</span><span class="pl-c1">/</span><span class="pl-s1">envs</span><span class="pl-c1">/</span><span class="pl-s1">pytorch</span><span class="pl-c1">/</span><span class="pl-s1">lib</span><span class="pl-c1">/</span><span class="pl-s1">python3</span>.<span class="pl-c1">8</span><span class="pl-c1">/</span><span class="pl-s1">site</span><span class="pl-c1">-</span><span class="pl-s1">packages</span><span class="pl-c1">/</span><span class="pl-s1">torch</span><span class="pl-c1">/</span><span class="pl-s1">autograd</span><span class="pl-c1">/</span><span class="pl-s1">__init__</span>.<span class="pl-s1">py</span> <span class="pl-c1">in</span> <span class="pl-en">backward</span>(<span class="pl-s1">tensors</span>, <span class="pl-s1">grad_tensors</span>, <span class="pl-s1">retain_graph</span>, <span class="pl-s1">create_graph</span>, <span class="pl-s1">grad_variables</span>)
     <span class="pl-c1">96</span>         <span class="pl-s1">retain_graph</span> <span class="pl-c1">=</span> <span class="pl-s1">create_graph</span>
     <span class="pl-c1">97</span> 
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">98</span>     <span class="pl-v">Variable</span>.<span class="pl-s1">_execution_engine</span>.<span class="pl-en">run_backward</span>(
     <span class="pl-c1">99</span>         <span class="pl-s1">tensors</span>, <span class="pl-s1">grad_tensors</span>, <span class="pl-s1">retain_graph</span>, <span class="pl-s1">create_graph</span>,
    <span class="pl-c1">100</span>         <span class="pl-s1">allow_unreachable</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)  <span class="pl-c"># allow_unreachable flag</span>

<span class="pl-v">RuntimeError</span>: <span class="pl-v">Function</span> <span class="pl-s">'SqrtBackward'</span> <span class="pl-s1">returned</span> <span class="pl-s1">nan</span> <span class="pl-s1">values</span> <span class="pl-c1">in</span> <span class="pl-s1">its</span> <span class="pl-c1">0</span><span class="pl-s1">th</span> <span class="pl-s1">output</span>.</pre></div>
<h2 dir="auto">Expected behavior</h2>
<p dir="auto">Training proceeds normally without nan gradients due to torch.cdist</p>
<h2 dir="auto">Environment</h2>
<ul dir="auto">
<li>PyTorch Version (e.g., 1.0): Pytorch 1.5</li>
<li>OS (e.g., Linux): Ubuntu 18.04 LTS</li>
<li>How you installed PyTorch (<code class="notranslate">conda</code>, <code class="notranslate">pip</code>, source): conda install pytorch torchvision cudatoolkit=10.1 -c pytorch</li>
<li>Python version: 3.8.2</li>
<li>CUDA/cuDNN version: 10.1/7.6</li>
<li>GPU models and configuration: RTX 2080 Ti/Titan Xp (tried both)</li>
<li>Any other relevant information: I believe things work in Pytorch 1.4</li>
</ul>
<p dir="auto">cc <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/ezyang/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/ezyang">@ezyang</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/gchanan/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/gchanan">@gchanan</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/zou3519/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/zou3519">@zou3519</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/fmassa/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/fmassa">@fmassa</a></p>