Update cvt.py

Author: fffffgggg54

timm/models/cvt.py

@@ -224,7 +224,7 @@ def __init__(
         qk_norm: bool = False,
         attn_drop: float = 0.,
         proj_drop: float = 0.,
-        input_norm_layer: nn.Module = partial(LayerNorm2d, eps=1e-5),
+        input_norm_layer: nn.Module = partial(LayerNorm, eps=1e-5),
         norm_layer: nn.Module = partial(LayerNorm, eps=1e-5),
         init_values: Optional[float] = None,
         drop_path: float = 0.,
@@ -326,7 +326,7 @@ def __init__(
         qk_norm: bool = False,
         attn_drop: float = 0.,
         proj_drop: float = 0.,
-        input_norm_layer = LayerNorm2d,
+        input_norm_layer = LayerNorm,
         norm_layer: nn.Module = LayerNorm,
         init_values: Optional[float] = None,
         drop_path_rates: List[float] = [0.],
@@ -417,7 +417,7 @@ def __init__(
         qk_norm: bool = False,
         attn_drop: float = 0.,
         proj_drop: float = 0.,
-        input_norm_layer = partial(LayerNorm2d, eps=1e-5),
+        input_norm_layer = partial(LayerNorm, eps=1e-5),
         norm_layer: nn.Module = partial(LayerNorm, eps=1e-5),
         init_values: Optional[float] = None,
         drop_path_rate: float = 0.,