Upload model

modeling.py

@@ -18,7 +18,7 @@ from transformers.utils.backbone_utils import load_backbone
 from .configuration import LSPDetrConfig, STAConfig
 
 
-flex_attention = torch.compile(flex_attention, dynamic=True)
+flex_attention = torch.compile(flex_attention, dynamic=False)
 
 
 def init_freqs(head_dim: int, num_heads: int, pos_dim: int, theta: float) -> Tensor:
@@ -106,6 +106,7 @@ class CayleySTRING(nn.Module):
             positions ([b, n, pos_dim]): Positions tensor.
         """
         # Compute (I + S)^-1 @ x
+        print(self.training)
         if self.training:
             # Use linalg.solve during training for numerical stability.
             y = torch.linalg.solve(
@@ -518,19 +519,7 @@ class LSPTransformer(nn.Module):
             "absolute_points": relative_to_absolute_pos(
                 ref_points, self.query_block_size, self.query_block_size
             ).flatten(1, 2),
-            "aux_outputs": [
-                {
-                    "logits": a,
-                    "points": b,
-                    "radial_distances": c,
-                }
-                for a, b, c in zip(
-                    logits_list[:-1],
-                    ref_points_list[:-1],
-                    radial_distances_list[:-1],
-                    strict=True,
-                )
-            ],
+            "embeddings": tgt.flatten(1, 2),
         }