Enable more test_torchscript (#16679)

* update _create_and_check_torchscript

* Enable test_torchscript

* clear_class_registry
Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>

tests/speech_to_text/test_modeling_speech_to_text.py

@@ -273,7 +273,6 @@ class Speech2TextModelTest(ModelTesterMixin, GenerationTesterMixin, unittest.Tes
     is_encoder_decoder = True
     test_pruning = False
     test_missing_keys = False
-    test_torchscript = True
 
     input_name = "input_features"
 

tests/squeezebert/test_modeling_squeezebert.py

@@ -229,7 +229,6 @@ class SqueezeBertModelTest(ModelTesterMixin, unittest.TestCase):
         else None
     )
     test_pruning = False
-    test_torchscript = True
     test_resize_embeddings = True
     test_head_masking = False
 

tests/swin/test_modeling_swin.py

@@ -177,7 +177,6 @@ class SwinModelTest(ModelTesterMixin, unittest.TestCase):
     )
 
     test_pruning = False
-    test_torchscript = False
     test_resize_embeddings = False
     test_head_masking = False
 

tests/t5/test_modeling_t5.py

@@ -512,7 +512,6 @@ class T5ModelTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCase):
     fx_compatible = True
     all_parallelizable_model_classes = (T5Model, T5ForConditionalGeneration) if is_torch_available() else ()
     test_pruning = False
-    test_torchscript = True
     test_resize_embeddings = True
     test_model_parallel = True
     is_encoder_decoder = True
@@ -777,7 +776,6 @@ class T5EncoderOnlyModelTester:
 class T5EncoderOnlyModelTest(ModelTesterMixin, unittest.TestCase):
     all_model_classes = (T5EncoderModel,) if is_torch_available() else ()
     test_pruning = False
-    test_torchscript = True
     test_resize_embeddings = False
     test_model_parallel = True
     all_parallelizable_model_classes = (T5EncoderModel,) if is_torch_available() else ()

tests/tapas/test_modeling_tapas.py

@@ -422,7 +422,6 @@ class TapasModelTest(ModelTesterMixin, unittest.TestCase):
         else None
     )
     test_pruning = False
-    test_torchscript = False
     test_resize_embeddings = True
     test_head_masking = False
 

tests/test_modeling_common.py

@@ -617,19 +617,21 @@ class ModelTesterMixin:
             model.eval()
             inputs = self._prepare_for_class(inputs_dict, model_class)
 
+            main_input_name = model_class.main_input_name
+
             try:
                 if model.config.is_encoder_decoder:
                     model.config.use_cache = False  # FSTM still requires this hack -> FSTM should probably be refactored similar to BART afterward
-                    input_ids = inputs["input_ids"]
+                    main_input = inputs[main_input_name]
                     attention_mask = inputs["attention_mask"]
                     decoder_input_ids = inputs["decoder_input_ids"]
                     decoder_attention_mask = inputs["decoder_attention_mask"]
                     traced_model = torch.jit.trace(
-                        model, (input_ids, attention_mask, decoder_input_ids, decoder_attention_mask)
+                        model, (main_input, attention_mask, decoder_input_ids, decoder_attention_mask)
                     )
                 else:
-                    input_ids = inputs["input_ids"]
-                    traced_model = torch.jit.trace(model, input_ids)
+                    main_input = inputs[main_input_name]
+                    traced_model = torch.jit.trace(model, main_input)
             except RuntimeError:
                 self.fail("Couldn't trace module.")
 

tests/transfo_xl/test_modeling_transfo_xl.py

@@ -238,7 +238,6 @@ class TransfoXLModelTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestC
     )
     all_generative_model_classes = (TransfoXLLMHeadModel,) if is_torch_available() else ()
     test_pruning = False
-    test_torchscript = False
     test_resize_embeddings = True
     test_mismatched_shapes = False
 

tests/unispeech/test_modeling_unispeech.py

@@ -305,7 +305,6 @@ class UniSpeechRobustModelTest(ModelTesterMixin, unittest.TestCase):
     )
     test_pruning = False
     test_headmasking = False
-    test_torchscript = False
 
     def setUp(self):
         self.model_tester = UniSpeechModelTester(

tests/van/test_modeling_van.py

@@ -124,7 +124,6 @@ class VanModelTest(ModelTesterMixin, unittest.TestCase):
     all_model_classes = (VanModel, VanForImageClassification) if is_torch_available() else ()
 
     test_pruning = False
-    test_torchscript = False
     test_resize_embeddings = False
     test_head_masking = False
     has_attentions = False

tests/vit/test_modeling_vit.py

@@ -158,7 +158,6 @@ class ViTModelTest(ModelTesterMixin, unittest.TestCase):
     )
 
     test_pruning = False
-    test_torchscript = False
     test_resize_embeddings = False
     test_head_masking = False
 

tests/wav2vec2/test_modeling_wav2vec2.py

@@ -413,7 +413,6 @@ class Wav2Vec2ModelTest(ModelTesterMixin, unittest.TestCase):
     )
     test_pruning = False
     test_headmasking = False
-    test_torchscript = False
 
     def setUp(self):
         self.model_tester = Wav2Vec2ModelTester(self)
@@ -652,7 +651,6 @@ class Wav2Vec2RobustModelTest(ModelTesterMixin, unittest.TestCase):
     )
     test_pruning = False
     test_headmasking = False
-    test_torchscript = False
 
     def setUp(self):
         self.model_tester = Wav2Vec2ModelTester(

tests/wavlm/test_modeling_wavlm.py

@@ -316,7 +316,6 @@ class WavLMModelTest(ModelTesterMixin, unittest.TestCase):
     )
     test_pruning = False
     test_headmasking = False
-    test_torchscript = False
 
     def setUp(self):
         self.model_tester = WavLMModelTester(self)