diff --git a/minerva/utils/position_embedding.py b/minerva/utils/position_embedding.py
index bc361ca..bf793d1 100644
--- a/minerva/utils/position_embedding.py
+++ b/minerva/utils/position_embedding.py
@@ -19,8 +19,8 @@ def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
     return:
     pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
     """
-    grid_h = np.arange(grid_size, dtype=np.float32)
-    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid_h = np.arange(grid_size, dtype=float)
+    grid_w = np.arange(grid_size, dtype=float)
     grid = np.meshgrid(grid_w, grid_h)  # here w goes first
     grid = np.stack(grid, axis=0)
 
@@ -49,7 +49,7 @@ def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
     out: (M, D)
     """
     assert embed_dim % 2 == 0
-    omega = np.arange(embed_dim // 2, dtype=np.float)
+    omega = np.arange(embed_dim // 2, dtype=float)
     omega /= embed_dim / 2.0
     omega = 1.0 / 10000**omega  # (D/2,)
 
diff --git a/tests/models/nets/test_setr.py b/tests/models/nets/test_setr.py
index b28a048..4909691 100644
--- a/tests/models/nets/test_setr.py
+++ b/tests/models/nets/test_setr.py
@@ -5,10 +5,10 @@
 
 
 def test_wisenet_loss():
-    model = SETR_PUP()
+    model = SETR_PUP(image_size=16)
     batch_size = 2
-    x = torch.rand(2, 3, 512, 512)
-    mask = torch.rand(2, 1, 512, 512).long()
+    x = torch.rand(2, 3, 16, 16)
+    mask = torch.rand(2, 1, 16, 16).long()
 
     # Do the training step
     loss = model.training_step((x, mask), 0).item()
@@ -17,11 +17,11 @@ def test_wisenet_loss():
 
 
 def test_wisenet_predict():
-    model = SETR_PUP()
+    model = SETR_PUP(image_size=16)
     batch_size = 2
-    mask_shape = (batch_size, 1000, 512, 512)  # (2, 1, 500, 500)
-    x = torch.rand(2, 3, 512, 512)
-    mask = torch.rand(2, 1, 512, 512).long()
+    mask_shape = (batch_size, 1000, 16, 16)  # (2, 1, 500, 500)
+    x = torch.rand(2, 3, 16, 16)
+    mask = torch.rand(2, 1, 16, 16).long()
 
     # Do the prediction step
     preds = model.predict_step((x, mask), 0)
diff --git a/tests/models/nets/test_sfm.py b/tests/models/nets/test_sfm.py
index 8f1f8ac..46f6ed2 100644
--- a/tests/models/nets/test_sfm.py
+++ b/tests/models/nets/test_sfm.py
@@ -20,7 +20,7 @@
 ]
 
 
-# @pytest.mark.parametrize("model_cls,img_size", test_models)
+@pytest.mark.parametrize("model_cls,img_size", test_models)
 def test_sfm_pretrain_forward(model_cls, img_size):
     # Test the class instantiation
     model = model_cls(img_size=img_size, in_chans=1, norm_pix_loss=False)