It’s simple, a != a will give you a ByteTensor, indicating the positions of NaNs
>>> b

    nan     nan -0.8395
    nan     nan     nan
-1.7921     nan  0.1864
[torch.FloatTensor of size 3x3]

>>> b != b

 1  1  0
 1  1  1
 0  1  0
[torch.ByteTensor of size 3x3]

you can use b[b != b] = 0 to set all NaNs to zero.