test: update tests

Author: okunator

cellseg_models_pytorch/decoders/tests/test_decoders.py

@@ -4,12 +4,14 @@
 from cellseg_models_pytorch.decoders import Decoder
 
 
-@pytest.mark.parametrize("long_skip", ["unet", "unetpp", "unet3p", "unet3p-lite"])
+@pytest.mark.parametrize(
+    "long_skip", ["unet", "unetpp", "unet3p", "unet3p-lite", "cross-attn"]
+)
 @pytest.mark.parametrize("merge_policy", ["cat", "sum"])
 @pytest.mark.parametrize("use_conv", [True, False])
 @pytest.mark.parametrize("use_tr", [True, False])
 def test_decoder_fwdbwd(long_skip, merge_policy, use_conv, use_tr):
-    enc_channels = (64, 32, 64, 32)
+    enc_channels = (32, 32, 32, 32)
     out_dims = [32 // 2**i for i in range(4)][::-1]
 
     decoder1_kwargs = {"merge_policy": merge_policy}
@@ -40,8 +42,8 @@ def test_decoder_fwdbwd(long_skip, merge_policy, use_conv, use_tr):
     decoder = Decoder(
         enc_channels=enc_channels,
         out_channels=(32, 32, 32, 32),
-        n_layers=n_layers,
-        n_blocks=n_blocks,
+        n_conv_layers=n_layers,
+        n_conv_blocks=n_blocks,
         n_transformers=n_tr_layers,
         n_transformer_blocks=n_tr_blocks,
         long_skip=long_skip,

cellseg_models_pytorch/models/tests/test_models.py

@@ -84,8 +84,8 @@ def test_multitaskunet_fwdbwd():
     m = MultiTaskUnet(
         decoders=("sem",),
         heads={"sem": {"sem": 3}},
-        n_layers={"sem": (1, 1, 1, 1)},
-        n_blocks={"sem": ((2,), (2,), (2,), (2,))},
+        n_conv_layers={"sem": (1, 1, 1, 1)},
+        n_conv_blocks={"sem": ((2,), (2,), (2,), (2,))},
         out_channels={"sem": (128, 64, 32, 16)},
         long_skips={"sem": "unet"},
         dec_params={"sem": None},

cellseg_models_pytorch/modules/tests/test_transformer.py

@@ -21,6 +21,7 @@ def test_transformer(block_type, computation_type):
         computation_types=(computation_type,),
         biases=(False,),
         dropouts=(0.0,),
+        layer_scales=(False,),
         slice_size=4,
         seq_len=H * W,
     )