test(modules): update transformer tests

Author: okunator

cellseg_models_pytorch/modules/tests/test_transformer.py

@@ -4,8 +4,9 @@
 from cellseg_models_pytorch.modules import Transformer2D
 
 
-@pytest.mark.parametrize("block_type", ["basic", "slice"])
-def test_transformer(block_type):
+@pytest.mark.parametrize("block_type", ["exact", "linformer"])
+@pytest.mark.parametrize("computation_type", ["basic", "slice"])
+def test_transformer(block_type, computation_type):
     in_channels = 64
     B = 4
     H = W = 32
@@ -17,9 +18,11 @@ def test_transformer(block_type):
         head_dim=32,
         n_blocks=1,
         block_types=(block_type,),
+        computation_types=(computation_type,),
         biases=(False,),
         dropouts=(0.0,),
         slice_size=4,
+        seq_len=H * W,
     )
 
     out = tr(x)