From 0317fe743300b831edf6b38840fe22f07af2a9aa Mon Sep 17 00:00:00 2001
From: Zeeshan Khan Suri <5270999+zshn25@users.noreply.github.com>
Date: Thu, 30 Nov 2023 17:58:18 +0100
Subject: [PATCH] added back fmha import

---
 dinov2/layers/block.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dinov2/layers/block.py b/dinov2/layers/block.py
index 68c6f45a9..638907474 100644
--- a/dinov2/layers/block.py
+++ b/dinov2/layers/block.py
@@ -27,7 +27,7 @@
 XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
 try:
     if XFORMERS_ENABLED:
-        from xformers.ops import scaled_index_add as _scaled_index_add, index_select_cat as _index_select_cat
+        from xformers.ops import fmha, scaled_index_add as _scaled_index_add, index_select_cat as _index_select_cat
 
         def scaled_index_add(input, index, source, scaling, alpha):
             is_proper_embed_dim = input.shape[-1] % 256 == 0