Update vision_transformer.py (facebookresearch#331)

Account for register tokens in get_intermediate_layers
wtwan · Dec 1, 2023 · 2302b6b · 2302b6b
1 parent da4b382
commit 2302b6b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dinov2/models/vision_transformer.py b/dinov2/models/vision_transformer.py
@@ -306,7 +306,7 @@ def get_intermediate_layers(
         if norm:
             outputs = [self.norm(out) for out in outputs]
         class_tokens = [out[:, 0] for out in outputs]
-        outputs = [out[:, 1:] for out in outputs]
+        outputs = [out[:, 1 + self.num_register_tokens:] for out in outputs]
         if reshape:
             B, _, w, h = x.shape
             outputs = [