Upload models/dinov2_encoder.py with huggingface_hub

models/dinov2_encoder.py

@@ -60,11 +60,9 @@ class Dinov2Encoder(nn.Module):
         """Forward through HF DINOv2, collecting per-layer features + final norm output."""
         features = {}
         model = self.encoder
-        batch_size = x.shape[0]
 
-        # HF DINOv2: patch_embeddings + CLS token + position embeddings
+        # HF DINOv2: embeddings (patch + CLS + pos) -> encoder layers -> norm
         embeddings = model.embeddings(x)
-        embeddings = model.layernorm_pre(embeddings)
 
         # Forward through each encoder layer
         for i, layer in enumerate(model.encoder.layer):
@@ -76,7 +74,7 @@ class Dinov2Encoder(nn.Module):
                 embeddings = outputs
             features[f'layer_{i}'] = embeddings[:, 1:]  # Remove CLS token
 
-        x_norm = model.layernorm(embeddings)
+        x_norm = model.norm(embeddings)
         return features, x_norm[:, 1:]
 
     def preprocess(self, lr: torch.Tensor) -> torch.Tensor: