fix patch

Author: xadupre

_unittests/ut_tasks/test_tasks_mask_generation.py

@@ -23,7 +23,7 @@ def test_mask_generation(self):
         model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
         model(**torch_deepcopy(inputs))
         model(**data["inputs2"])
-        with torch_export_patches(patch_transformers=True, verbose=10):
+        with torch_export_patches(patch_transformers=True, verbose=1):
             torch.export.export(
                 model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
             )

onnx_diagnostic/torch_export_patches/patches/patch_transformers.py

@@ -1300,14 +1300,19 @@ def sparse_prompt_embeddings_is_empty(output_tokens, sparse_prompt_embeddings):
         )
 
         # Run the transformer, image_positional_embedding are consumed
-        point_embedding, image_embeddings, attentions = self.transformer(
+        torch._check(point_embeddings.shape[0] != 0)
+        torch._check(point_embeddings.shape[1] != 0)
+        torch._check(point_embeddings.shape[2] != 0)
+        torch._check(point_embeddings.shape[3] != 0)
+        embeddings_attentions = self.transformer(
             point_embeddings=point_embeddings,
             image_embeddings=image_embeddings,
             image_positional_embeddings=image_positional_embeddings,
             attention_similarity=attention_similarity,
             target_embedding=target_embedding,
             output_attentions=output_attentions,
         )
+        point_embedding, image_embeddings = embeddings_attentions[:2]
         iou_token_out = torch.select(point_embedding, dim=2, index=0)
         mask_tokens_out = torch.narrow(
             point_embedding, dim=2, start=1, length=self.num_mask_tokens
@@ -1349,9 +1354,12 @@ def sparse_prompt_embeddings_is_empty(output_tokens, sparse_prompt_embeddings):
 
         outputs = (masks, iou_pred)
 
-        if output_attentions:
-            outputs = outputs + (attentions,)  # noqa: RUF005
+        if len(embeddings_attentions) == 2:
+            # transformers==4.54
+            return outputs
+
+        if output_attentions and len(embeddings_attentions) > 2:
+            outputs = outputs + (embeddings_attentions[2],)  # noqa: RUF005
         else:
             outputs = outputs + (None,)  # noqa: RUF005
-
         return outputs