match mm branch

tests/test_feature_extractor.py

@@ -12,7 +12,7 @@ def test_feature_extractor():
     device = "cuda" if torch.cuda.is_available() else "cpu"
     segmentation_types = ["none", "grid", "slic", "random", "stego"]
     feature_types = ["dino", "dinov2", "stego"]
-    backbone_types = ["vit_small", "vit_base", "vit_small_reg", "vit_base_reg"]
+    backbone_types = ["vit_small", "vit_base"] # "vit_small_reg", "vit_base_reg"]
 
     for seg_type, feat_type, back_type in itertools.product(segmentation_types, feature_types, backbone_types):
         if seg_type == "stego" and feat_type != "stego":
@@ -39,7 +39,7 @@ def test_feature_extractor():
 
         ax[0].imshow(transform(img).permute(0, 2, 3, 1)[0].cpu())
         ax[0].set_title("Image")
-        ax[1].imshow(seg.cpu(), cmap=plt.colormaps.get("inferno"))
+        ax[1].imshow(seg.cpu(), cmap=plt.colormaps.get("gray"))
         ax[1].set_title("Segmentation")
         plt.tight_layout()