Merge pull request #18 from Mr-Grin/fix-selected-index-k-out-of-range

fix selected index k out of range

Author: lucidrains

native_sparse_attention_pytorch/native_sparse_attention.py

@@ -675,6 +675,8 @@ def forward(
         fk = k
         fv = v
 
+        num_selected = min(num_selected, importance_scores.shape[-1] - 1)
+        
         if has_selected_kv_for_fine_attn:
 
             # get the top-n kv segments for fine attention