Issue with torch.layer_norm Function Returning NaN Values

Hi,
I’m encountering an issue with the torch.layer_norm function, where it returns a Tensor containing NaN values under certain inputs. I’m unsure if this behavior indicates a bug or if it’s expected.

Here’s a minimal code snippet that reproduces the problem:

import torch

features = torch.load('features.pt')
print(torch.isnan(features).any())
print(torch.isnan(torch.layer_norm(features, (64,))).any())

And here a colab notebook that reproduces the problem.

Please find the input tensor here.

Thank you for your help!