fix(modules): fix layer norm for chl-attn

Author: okunator

cellseg_models_pytorch/modules/attention_modules.py

@@ -257,7 +257,7 @@ def __init__(
             padding=0,
             bias=True,
         )
-        self.norm = Norm("ln", num_features=squeeze_channels)
+        self.norm = Norm("ln2d", num_features=squeeze_channels)
         self.act = Activation(activation)
         self.conv_excite = Conv(
             conv,

cellseg_models_pytorch/modules/tests/test_basemodules.py

@@ -13,7 +13,7 @@ def test_act_forward(activation):
     assert output.dtype == input.dtype
 
 
-@pytest.mark.parametrize("normalization", ["bn", "bcn", "gn", "ln", None])
+@pytest.mark.parametrize("normalization", ["bn", "bcn", "gn", "ln2d", None])
 def test_norm(normalization):
     norm = Norm(normalization, num_features=3)
     input = torch.rand([1, 3, 16, 16])