GroupNorm

import torch
import torch.nn as nn

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

# Input data: (batch_size, channels, height, width)
input = torch.randn(2, 6, 32, 32) # 2 images, 6 channels, 32 × 32 resolution, default data type: fp32

# Define GroupNorm (3 groups, 2 channels in each group).
group_norm = nn.GroupNorm(num_groups=3, num_channels=6) # 6 channels are divided into 3 groups. The default data type is fp32.

# Forward computation
output = group_norm(input)

# Verification: normalization within each group
print("GroupNorm output shape: ", output.shape # Retain the original shape.