Question about computational improvement of LayerNorm2d

Thank you for your excellent work.

I would like to inquire whether a certain modification could potentially enhance the speed of your model's computation.
The `LayerNorm2d` involves the operation `[(x - u) / s] * w + b`.
If we perform this operation by combining the `/s` and `*w` into `*(w/s)`, it will require only 3 computations of `x.shape` instead of 4, thereby reducing the time usage and memory usage during the computation.

The following is my modification
```python
class LayerNorm2d(nn.Module):
def __init__(self, num_channels: int, eps: float = 1e-6) -> None:
super().__init__()
self.weight = nn.Parameter(torch.ones(num_channels))
self.bias = nn.Parameter(torch.zeros(num_channels))
self.eps = eps

def forward(self, x: torch.Tensor) -> torch.Tensor:
mean = x.mean(1, keepdim=True)
x = x - mean
variance = x.pow(2).mean(1, keepdim=True)
scaling = self.weight[:, None, None] / torch.sqrt(variance + self.eps)
x = scaling * x + self.bias[:, None, None]
return x
```

Start by locating the LayerNorm2d implementation and compare it with the proposed forward method in this issue. Benchmark both versions for computation time and memory usage, then validate that the output remains equivalent before deciding whether the change is worthwhile.

Written by the indexing model from the issue text.