From b8e789ce84c300afd4a7a5d91545e68af13ffc19 Mon Sep 17 00:00:00 2001
From: qasfb <36480216+qasfb@users.noreply.github.com>
Date: Fri, 1 Dec 2023 18:00:30 +0100
Subject: [PATCH] Update vision_transformer.py
Account for register tokens in get_intermediate_layers
---
dinov2/models/vision_transformer.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/dinov2/models/vision_transformer.py b/dinov2/models/vision_transformer.py
index c8c3ec2..4926108 100644
--- a/dinov2/models/vision_transformer.py
+++ b/dinov2/models/vision_transformer.py
@@ -306,7 +306,7 @@ class DinoVisionTransformer(nn.Module):
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 = [
--
GitLab