diff --git a/pyretri/index/re_ranker/re_ranker_impl/query_expansion.py b/pyretri/index/re_ranker/re_ranker_impl/query_expansion.py
index 0f02982..b19e65e 100644
--- a/pyretri/index/re_ranker/re_ranker_impl/query_expansion.py
+++ b/pyretri/index/re_ranker/re_ranker_impl/query_expansion.py
@@ -59,9 +59,9 @@ class QE(ReRankerBase):
             query_fea = requery_fea
             dis = self._cal_dis(query_fea, gallery_fea)
 
-        if kr is None:
-            sorted_index = torch.argsort(dis, dim=1)
-        else:
-            sorted_index = kr(query_fea, gallery_fea, dis)
+            if kr is None:
+                sorted_index = torch.argsort(dis, dim=1)
+            else:
+                sorted_index = kr(query_fea, gallery_fea, dis)
 
         return sorted_index