Gradient of a weight matrix returns NAN with respect to learnable parameter

username: ABP
cooked: <p>I have this function which i use to evaluate weights for my integration</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/1/6130ca1ee7e23292748a4e65e8704e55470e43da.png" data-download-href="https://discuss.pytorch.org/uploads/default/6130ca1ee7e23292748a4e65e8704e55470e43da" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/6/1/6130ca1ee7e23292748a4e65e8704e55470e43da.png" alt="image" data-base62-sha1="dRMQqYIGo0mgVham4yghrfxzHwm" width="690" height="228" data-dominant-color="27292A"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">868×288 16.8 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I am trying to use it as pat of the computational graph. I found out that the gradient of the returned matrix w with respect to a is returned as Nan. Hence the optimiser takes NaN value for the learnable parameter a in a PINN routine.How do I solve the NaN problem?</p>

username: albanD
cooked: <p>Hey!</p>
<p>It’s a bit hard to say with this but my best guess is that you’re using your “mask” to remove unvalid values. Unfortunately that means you fall into the “non differentiable” part of the computation and that can play tricks for you.</p>
<p>I would suggest:</p>
<ul>
<li>If any part of your computation is generating invalid values, you should mask them out BEFORE they’re computed in the forward. If you do it after, then in the backward, these invalid values will happen after the masking and thus lead to NaN</li>
<li>You can use anomaly mode in autograd to knew exactly which op generated NaNs and get a better sense of what’s happening in the backward: <a href="https://pytorch.org/docs/stable/autograd.html#torch.autograd.detect_anomaly" class="inline-onebox">Automatic differentiation package - torch.autograd — PyTorch 2.7 documentation</a></li>
</ul>

username: ABP
cooked: <aside class="quote no-group" data-username="albanD" data-post="2" data-topic="219604">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>Unfortunately that means you fall into the “non differentiable” part of the computation and that can play tricks for you</p>
</blockquote>
</aside>
<p>Its quite amusing when i define the weights using  for loop , the gradient computation seems to flow properly. But i am trying to vectorize a Finite difference scheme for a numerical solver. I checked the computation in forward pass , it seems to work fine. I did check the anomaly mode in autograd for the backward pass. And the highlighted part of the code is the culprit for the NaN gradient. I tried using a custom gradient evaluation for the weight matrix( using a finite difference for gradient evaluation) and linked it to the chain rule of autodiff. Prediction seems of as gradient of learnable parameter is very small</p>