fix tests

Author: AlbertDominguez

tests/test_decoders.py

@@ -9,7 +9,6 @@ def test_decoder_2d():
         dim=64,
         num_layers=3,
         heads=2,
-        use_rotary_embed=True,
         dropout=0.0,
     )
     B = 2
@@ -27,7 +26,6 @@ def test_decoder_3d():
         dim=64,
         num_layers=3,
         heads=2,
-        use_rotary_embed=True,
         dropout=0.0,
     )
     B = 2
@@ -44,7 +42,6 @@ def test_decoder_4d():
         dim=128,
         num_layers=3,
         heads=2,
-        use_rotary_embed=True,
         dropout=0.0,
     )
     B = 2

tests/test_mae.py

@@ -63,7 +63,7 @@ def __getitem__(self, idx: int):
 
 @pytest.mark.parametrize("in_channels", [1, 3])
 @pytest.mark.parametrize("ndim", [2, 3, 4])
-def test_mae_fwd(input_space: str, in_channels: int, ndim: int):
+def test_mae_fwd(in_channels: int, ndim: int):
     if ndim not in (2, 3, 4):
         raise ValueError("ndim must be either 2, 3, or 4.")
     ActualCls = NDIM_TO_MAE_CLS[ndim]
@@ -168,4 +168,4 @@ def test_mae_io(ndim: int):
     )
 
 if __name__ == "__main__":
-    test_mae_fwd("real", 1, 4)
+    test_mae_fwd("real", 1, 4)