From 8cbdd8f72e4fba8ef66ac936945e87f0fdc41943 Mon Sep 17 00:00:00 2001
From: Talmaj Marinc <talmaj@comfy.org>
Date: Mon, 8 Jun 2026 12:46:58 +0200
Subject: [PATCH] Use optimized attention in Dino2AttentionBlock.

---
 comfy/image_encoders/dino2.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/comfy/image_encoders/dino2.py b/comfy/image_encoders/dino2.py
index 2f51693be..40f569ad8 100644
--- a/comfy/image_encoders/dino2.py
+++ b/comfy/image_encoders/dino2.py
@@ -53,8 +53,7 @@ class Dino2AttentionBlock(torch.nn.Module):
         if rope is not None and pos is not None:
             q = rope(q, pos)
             k = rope(k, pos)
-        out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
-        out = out.transpose(1, 2).reshape(B, N, C)
+        out = optimized_attention(q, k, v, h, mask=mask, skip_reshape=True)
         return self.output(out)