Fixes a dtype mismatch in LayerNorm2D

Author: pranavm-nvidia

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

@@ -186,7 +186,11 @@ def __init__(self, num_channels: int, eps: float = 1e-6) -> None:
 
     def forward(self, x: tp.Tensor) -> tp.Tensor:
         x = tp.permute(x, (0, 2, 3, 1))
+        # LayerNorm is always done in float32:
+        original_dtype = x.dtype
+        x = tp.cast(x, tp.float32)
         x = super().forward(x)
+        x = tp.cast(x, original_dtype)
         return tp.permute(x, (0, 3, 1, 2))