Adapt checking logic

src/adapters/methods/reft.py

@@ -71,7 +71,11 @@ def _gather_adapted_states(self, hidden_states: torch.Tensor):
         # if cached indexing matrices are computed for different hidden_states size -> recompute
         cache_invalidated = False
         if hasattr(context, "pref_idx") and hasattr(context, "suff_idx"):
-            cache_invalidated = context.suff_idx.size(1) != seq_len
+            cache_invalidated = (
+                torch.max(context.suff_idx) >= seq_len  # indices out of bounds
+                or bsz != context.suff_idx.size(0)  # batch size mismatch
+                or ddim != context.suff_idx.size(2)  # hidden size mismatch
+            )
 
         # no cached indexing matrices available -> compute now
         if not hasattr(context, "pref_idx") and not hasattr(context, "suff_idx") or cache_invalidated: