Update LayerNorm2D to use TRT normalization API

Author: akhilg-nv

tripy/examples/segment-anything-model-v2/sam2/modeling/sam2_utils.py

@@ -188,18 +188,14 @@ def __init__(self, num_channels: int, eps: float = 1e-6) -> None:
         self.weight = DefaultParameter((num_channels,), tp.float32)
         self.bias = DefaultParameter((num_channels,), tp.float32)
         self.eps = eps
+        self.num_channels = num_channels
 
     def forward(self, x: tp.Tensor) -> tp.Tensor:
-        original_dtype = x.dtype
-        x = tp.cast(x, tp.float32)
-        u = tp.mean(x, dim=1, keepdim=True)
-        s = tp.mean((x - u) ** 2, dim=1, keepdim=True)
-        x = (x - u) / tp.sqrt(s + self.eps)
-        w = tp.unsqueeze(tp.unsqueeze(self.weight, 1), 2)
-        b = tp.unsqueeze(tp.unsqueeze(self.bias, 1), 2)
-        x = w * x + b
-        x = tp.cast(x, original_dtype)
-        return x
+        x_perm = tp.permute(x, (0, 2, 3, 1))
+        ln = tp.LayerNorm(self.num_channels, dtype=x.dtype, eps=self.eps)
+        ln.weight = self.weight
+        ln.bias = self.bias
+        return tp.permute(ln(x_perm), (0, 3, 1, 2))
 
 
 def get_activation_fn(activation):