def __init__(
            self,
            num_channels: int,
            eps: float = 1e-12,
    ):
        """Uses GroupNorm implementation with group=1 for speed."""
        super().__init__()
        # we use GroupNorm to implement this efficiently and fast.
        self.layer_norm = torch.nn.GroupNorm(1, num_channels=num_channels, eps=eps)
Uses GroupNorm implementation with group=1 for speed.
analyze/convnexts4nd/src/models/baselines/ckconv.py