Update ml_decoder.py

huggingface · Dec 29, 2024 · 5cac39c · 5cac39c
1 parent fbeed51
commit 5cac39c
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/timm/layers/ml_decoder.py b/timm/layers/ml_decoder.py
@@ -436,14 +436,14 @@ def _resolve_query(self, q):
         if q is not None:
             return q
         if not self.have_class_embed:
-            return self.query_embed
+            return self.query_embed.weight
         else:
             if self.class_embed_merge == 'add':
-                return self.query_embed + self.class_embed
+                return self.query_embed.weight + self.class_embed.weight
             elif self.class_embed_merge == 'concat':
                 return torch.cat([x.weight for x in [self.query_embed, self.class_embed] if x is not None], dim=1)
             else:
-                return self.class_embed
+                return self.class_embed.weight
 
     def forward(self, x, q=None):
         # BCHW to BNC