Division by zero or reaching infinity somewhere(not limited to these cases) can cause Nan.
I don’t think these two can cause Nan. These 2 functions are generally used Spatial Transformer Networks. affine_grid apply linear transformation like rotating, translating, scaling(I think sheering is also part of the transformation but I’m not sure) to the input image.  grid_sample uses the coordinates generated by affine_grid to perform a bi-linear interpolation on the original image to form a result image. https://pytorch.org/tutorials/intermediate/spatial_transformer_tutorial.html 2

Ur code is quite long   so I only took a glance. My suggestion is to check the part where u used .exp(), I had expierenced Nan using expontenial operations tho it may not be the case for u. If that’s the case, then the Nan in gradients can also be explained. U can also use tc.any(tc.isnan(INPUT)) to see where did Nan show up.

Hope this helps