Why this operation results in a tensor full of 'nan'?
err = reduce((mu_students - teacher_pred)**2, 'b h w vec -&gt; b h w', 'sum')
where mu_students and teacher_pred are 2 tensors with size (1,14,14,256).
The two tensors do not contain nan values before the reduce.
Moreover, this exact same operation work fine with the previous layer of the network, when mu_students and teacher_pred have size = (1,28,28,128). This error presents itself only on one kind of images (but it shouldn't be related to size, because images are always resized to 224)

Could you provide more about the function reduce and the input tensors?
By assuming that;

(1) reduce is from einops and written as
from einops import reduce as reduce

(2) mu_students and teacher_pred is random tensors,


when I try this code on Google Colab, it works well without NaN.
!pip install einops
from einops import reduce as reduce
mu_students = torch.rand(1,14,14,256)
teacher_pred = torch.rand(1,14,14,256)
err = reduce((mu_students - teacher_pred)**2, 'b h w vec -&gt; b h w', 'sum')

The values are, for example,
# mu_students
tensor([[[[0.4079, 0.5835, 0.6807,  ..., 0.6041, 0.7366, 0.9291],
          [0.0338, 0.9161, 0.7018,  ..., 0.6035, 0.1816, 0.4059],
          [0.5949, 0.9535, 0.1460,  ..., 0.4049, 0.5120, 0.0734], ...

# teacher_pred
tensor([[[[0.7493, 0.2193, 0.6465,  ..., 0.6262, 0.0270, 0.5532],
          [0.5343, 0.9384, 0.9916,  ..., 0.2127, 0.0370, 0.6322],
          [0.3568, 0.7474, 0.5562,  ..., 0.0589, 0.1356, 0.6062], ...
# err
tensor([[[42.6367, 36.5668, 42.7598, 37.0643, 40.7744, 43.6076, 45.5338,
          45.3421, 40.0736, 45.2035, 43.3516, 40.8768, 39.1142, 40.8736],
         [50.8344, 40.1921, 40.7194, 44.3139, 37.7520, 39.2365, ...

I wonder if this also works in your environment.
Do you use specific tensors as mu_students or teacher_pred?
