From e2910cb4589cab3db6b94d44576262b021a1d885 Mon Sep 17 00:00:00 2001
From: YannAhlgrim <yann@ahlgrim.net>
Date: Sun, 12 Apr 2026 11:19:06 +0200
Subject: [PATCH] Header for Linear Probing

---
 src/models/head.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)
 create mode 100644 src/models/head.py

diff --git a/src/models/head.py b/src/models/head.py
new file mode 100644
index 0000000..f846486
--- /dev/null
+++ b/src/models/head.py
@@ -0,0 +1,22 @@
+import torch
+import torch.nn as nn
+
+
+class ViTClassifier(nn.Module):
+    def __init__(self, encoder, num_classes, embed_dim):
+        super().__init__()
+        self.encoder = encoder
+        self.head = nn.Linear(embed_dim, num_classes)
+
+        nn.init.trunc_normal_(self.head.weight, std=0.01)
+        nn.init.zeros_(self.head.bias)
+
+    def forward(self, x):
+        # ViT -> (B, N, D)
+        features = self.encoder(x)
+
+        # Average Pool -> (B, D)
+        avg_embed = features.mean(dim=1)
+
+        logits = self.head(avg_embed)
+        return logits