LayerNorm

import torch
import torch.nn as nn

# Enable KuDNN.
torch._C._set_kdnn_enabled(True)

# Input data: (batch_size, seq_len, hidden_dim)
input = torch.randn(2, 5, 10) # 2 samples, 5 time steps, 10 feature dimensions, fp32 as the default data type

# Define LayerNorm (normalizing the last dimension hidden_dim=10).
layer_norm = nn.LayerNorm(10)  # Or normalized_shape=[5,10] to normalize the last two dimensions. The default data type is fp32.

# Forward computation
output = layer_norm(input)

# Verification: After normalization, the mean value of the last dimension is close to 0, and the variance is close to 1.
print("LayerNorm output mean:", output.mean(dim=-1)) # Should be close to 0.
print("LayerNorm output variance:", output.var(dim=-1, unbiased=False)) # Should be close to 1.