Fix UnboundLocalError in SuperPoint forward under torch.compile (#49409)

* fix bug for superpoint under compile mode

Signed-off-by: kaixuanliu <kaixuan.liu@intel.com>

* Update tests/models/superpoint/test_modeling_superpoint.py

Co-authored-by: guarin <43336610+guarin@users.noreply.github.com>

* update

Signed-off-by: kaixuanliu <kaixuan.liu@intel.com>

* update

Signed-off-by: kaixuanliu <kaixuan.liu@intel.com>

---------

Signed-off-by: kaixuanliu <kaixuan.liu@intel.com>
Co-authored-by: guarin <43336610+guarin@users.noreply.github.com>

Author: kaixuanliu

src/transformers/models/superpoint/modeling_superpoint.py

@@ -421,15 +421,15 @@ def forward(
         last_hidden_state = encoder_outputs[0]
 
         list_keypoints_scores = [
-            self.keypoint_decoder(last_hidden_state[None, ...]) for last_hidden_state in last_hidden_state
+            self.keypoint_decoder(image_hidden_state[None, ...]) for image_hidden_state in last_hidden_state
         ]
 
         list_keypoints = [keypoints_scores[0] for keypoints_scores in list_keypoints_scores]
         list_scores = [keypoints_scores[1] for keypoints_scores in list_keypoints_scores]
 
         list_descriptors = [
-            self.descriptor_decoder(last_hidden_state[None, ...], keypoints[None, ...])
-            for last_hidden_state, keypoints in zip(last_hidden_state, list_keypoints)
+            self.descriptor_decoder(image_hidden_state[None, ...], keypoints[None, ...])
+            for image_hidden_state, keypoints in zip(last_hidden_state, list_keypoints)
         ]
 
         maximum_num_keypoints = max(keypoints.shape[0] for keypoints in list_keypoints)

tests/models/superpoint/test_modeling_superpoint.py

@@ -15,8 +15,17 @@
 import unittest
 from functools import cached_property
 
+import pytest
+
 from transformers.models.superpoint.configuration_superpoint import SuperPointConfig
-from transformers.testing_utils import is_flaky, require_torch, require_vision, slow, torch_device
+from transformers.testing_utils import (
+    is_flaky,
+    require_torch,
+    require_torch_accelerator,
+    require_vision,
+    slow,
+    torch_device,
+)
 from transformers.utils import is_torch_available, is_vision_available
 
 from ...test_configuration_common import ConfigTester
@@ -239,6 +248,28 @@ def test_forward_labels_should_be_none(self):
                     model(**model_inputs)
                 self.assertEqual(ValueError, cm.exception.__class__)
 
+    @slow
+    @require_torch_accelerator
+    @pytest.mark.torch_compile_test
+    def test_compile_forward(self):
+        # Regression test for https://github.com/huggingface/transformers/issues/49561: the list comprehensions
+        # in the forward pass used to shadow `last_hidden_state`, which raised an `UnboundLocalError` in the
+        # resume function generated after a dynamo graph break.
+        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+        for model_class in self.all_model_classes:
+            model = model_class(config)
+            model.to(torch_device)
+            model.eval()
+
+            torch._dynamo.reset()
+            model.forward = torch.compile(model.forward)
+
+            with torch.no_grad():
+                outputs = model(**self._prepare_for_class(inputs_dict, model_class))
+
+            self.assertIsNotNone(outputs.keypoints)
+            self.assertIsNotNone(outputs.descriptors)
+
 
 def prepare_imgs():
     image1 = Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png")