Upload folder using huggingface_hub

modeling_gend.py

@@ -6,12 +6,15 @@ from transformers import PretrainedConfig, PreTrainedModel
 
 
 class LinearProbe(nn.Module):
-    def __init__(self, input_dim, num_classes, normalize_inputs=False, detach_classifier_inputs=False):
+    def __init__(self, input_dim, num_classes, normalize_inputs=False):
         super().__init__()
         self.linear = nn.Linear(input_dim, num_classes)
         self.normalize_inputs = normalize_inputs
 
     def forward(self, x: torch.Tensor, **kwargs):
+        if self.normalize_inputs:
+            x = F.normalize(x, p=2, dim=1)
+
         return self.linear(x)
 
 
@@ -50,10 +53,6 @@ class CLIPEncoder(nn.Module):
 
 class DINOEncoder(nn.Module):
     def __init__(self, model_name="facebook/dinov2-with-registers-base"):
-        """
-        See models in src/config.py
-        """
-
         super().__init__()
 
         from transformers import AutoImageProcessor, AutoModel, Dinov2Model, Dinov2WithRegistersModel
@@ -74,11 +73,7 @@ class DINOEncoder(nn.Module):
 
 
 class PerceptionEncoder(nn.Module):
-    def __init__(
-        self,
-        model_name="vit_pe_core_large_patch14_336",
-        img_size: None | int = None,
-    ):
+    def __init__(self, model_name="vit_pe_core_large_patch14_336", img_size: None | int = None):
         super().__init__()
 
         if img_size is not None: