Fix import of transformmodindex on nightly (#52)

stack-info: PR: #52, branch: drisspg/stack/1

Author: drisspg

attn_gym/utils.py

@@ -10,7 +10,11 @@
     _vmap_for_bhqkv,
     _ModificationType,
 )
-from torch._higher_order_ops.flex_attention import TransformGetItemToIndex
+# TODO This was moved on nightly, this enables 2.5 and 2.6 | we should remove this once 2.5 is no longer supported
+try:
+    from torch._dynamo._trace_wrapped_higher_order_op import TransformGetItemToIndex
+except ImportError:
+    from torch._higher_order_ops.flex_attention import TransformGetItemToIndex
 from contextlib import nullcontext
 
 Tensor = torch.Tensor