Fix torch.fx issue related to the new loss_kwargs keyword argument (huggingface#34380)

* Fix FX

* Unskip tests

Author: michaelbenayoun

src/transformers/utils/fx.py

@@ -1416,7 +1416,7 @@ def keys(self, obj: "Proxy") -> Any:
         your custom tracer.
         """
         attribute = HFAttribute(obj, "keys")()
-        if obj.node.target == "**kwargs":
+        if obj.node.target.startswith("**"):
             return attribute._metadata
         return attribute
 

tests/models/cohere/test_modeling_cohere.py

@@ -304,7 +304,6 @@ def test_model_various_embeddings(self):
             config_and_inputs[0].position_embedding_type = type
             self.model_tester.create_and_check_model(*config_and_inputs)
 
-    @unittest.skip(reason="PR #34283 made changes to the forward function.")
     def test_torch_fx_output_loss(self):
         super().test_torch_fx_output_loss()
 

tests/models/mistral/test_modeling_mistral.py

@@ -356,7 +356,6 @@ def test_model_various_embeddings(self):
             config_and_inputs[0].position_embedding_type = type
             self.model_tester.create_and_check_model(*config_and_inputs)
 
-    @unittest.skip(reason="PR #34283 made changes to the forward function.")
     def test_torch_fx_output_loss(self):
         super().test_torch_fx_output_loss()
 

tests/models/mixtral/test_modeling_mixtral.py

@@ -356,7 +356,6 @@ def test_model_various_embeddings(self):
             config_and_inputs[0].position_embedding_type = type
             self.model_tester.create_and_check_model(*config_and_inputs)
 
-    @unittest.skip(reason="PR #34283 made changes to the forward function.")
     def test_torch_fx_output_loss(self):
         super().test_torch_fx_output_loss()
 

tests/models/qwen2/test_modeling_qwen2.py

@@ -368,7 +368,6 @@ def test_model_various_embeddings(self):
             config_and_inputs[0].position_embedding_type = type
             self.model_tester.create_and_check_model(*config_and_inputs)
 
-    @unittest.skip(reason="PR #34283 made changes to the forward function.")
     def test_torch_fx_output_loss(self):
         super().test_torch_fx_output_loss()
 

tests/models/qwen2_moe/test_modeling_qwen2_moe.py

@@ -391,7 +391,6 @@ def test_model_various_embeddings(self):
             config_and_inputs[0].position_embedding_type = type
             self.model_tester.create_and_check_model(*config_and_inputs)
 
-    @unittest.skip(reason="PR #34283 made changes to the forward function.")
     def test_torch_fx_output_loss(self):
         super().test_torch_fx_output_loss()