Enable more test_torchscript (#16679)

ydshieh · web-flow · commit c04619ecf3ce · 2022-04-11T18:23:35.000+02:00
* update _create_and_check_torchscript

* Enable test_torchscript

* clear_class_registry

Co-authored-by: ydshieh &lt;ydshieh@users.noreply.github.com&gt;
diff --git a/tests/beit/test_modeling_beit.py b/tests/beit/test_modeling_beit.py
@@ -192,7 +192,6 @@ class BeitModelTest(ModelTesterMixin, unittest.TestCase):
     )
 
     test_pruning = False
-    test_torchscript = False
     test_resize_embeddings = False
     test_head_masking = False
 
diff --git a/tests/canine/test_modeling_canine.py b/tests/canine/test_modeling_canine.py
@@ -219,7 +219,6 @@ class CanineModelTest(ModelTesterMixin, unittest.TestCase):
         else ()
     )
 
-    test_torchscript = False
     test_mismatched_shapes = False
     test_resize_embeddings = False
     test_pruning = False
diff --git a/tests/clip/test_modeling_clip.py b/tests/clip/test_modeling_clip.py
@@ -151,7 +151,6 @@ class CLIPVisionModelTest(ModelTesterMixin, unittest.TestCase):
     all_model_classes = (CLIPVisionModel,) if is_torch_available() else ()
 
     test_pruning = False
-    test_torchscript = False
     test_resize_embeddings = False
     test_head_masking = False
 
diff --git a/tests/convnext/test_modeling_convnext.py b/tests/convnext/test_modeling_convnext.py
@@ -138,7 +138,6 @@ class ConvNextModelTest(ModelTesterMixin, unittest.TestCase):
     )
 
     test_pruning = False
-    test_torchscript = False
     test_resize_embeddings = False
     test_head_masking = False
     has_attentions = False
diff --git a/tests/ctrl/test_modeling_ctrl.py b/tests/ctrl/test_modeling_ctrl.py
@@ -174,7 +174,6 @@ class CTRLModelTest(ModelTesterMixin, GenerationTesterMixin, unittest.TestCase):
     all_model_classes = (CTRLModel, CTRLLMHeadModel, CTRLForSequenceClassification) if is_torch_available() else ()
     all_generative_model_classes = (CTRLLMHeadModel,) if is_torch_available() else ()
     test_pruning = True
-    test_torchscript = False
     test_resize_embeddings = False
     test_head_masking = False
 
diff --git a/tests/data2vec/test_modeling_data2vec_audio.py b/tests/data2vec/test_modeling_data2vec_audio.py
@@ -372,7 +372,6 @@ class Data2VecAudioModelTest(ModelTesterMixin, unittest.TestCase):
     )
     test_pruning = False
     test_headmasking = False
-    test_torchscript = False
 
     def setUp(self):
         self.model_tester = Data2VecAudioModelTester(self)
diff --git a/tests/decision_transformer/test_modeling_decision_transformer.py b/tests/decision_transformer/test_modeling_decision_transformer.py
@@ -148,6 +148,7 @@ class DecisionTransformerModelTest(ModelTesterMixin, GenerationTesterMixin, unit
     test_inputs_embeds = False
     test_model_common_attributes = False
     test_gradient_checkpointing = False
+    test_torchscript = False
 
     def setUp(self):
         self.model_tester = DecisionTransformerModelTester(self)
diff --git a/tests/deit/test_modeling_deit.py b/tests/deit/test_modeling_deit.py
@@ -169,7 +169,6 @@ class DeiTModelTest(ModelTesterMixin, unittest.TestCase):
     )
 
     test_pruning = False
-    test_torchscript = False
     test_resize_embeddings = False
     test_head_masking = False
 
diff --git a/tests/distilbert/test_modeling_distilbert.py b/tests/distilbert/test_modeling_distilbert.py
@@ -211,7 +211,6 @@ class DistilBertModelTest(ModelTesterMixin, unittest.TestCase):
     )
     fx_compatible = True
     test_pruning = True
-    test_torchscript = True
     test_resize_embeddings = True
     test_resize_position_embeddings = True
 
diff --git a/tests/dpt/test_modeling_dpt.py b/tests/dpt/test_modeling_dpt.py
@@ -154,7 +154,6 @@ class DPTModelTest(ModelTesterMixin, unittest.TestCase):
     all_model_classes = (DPTModel, DPTForDepthEstimation, DPTForSemanticSegmentation) if is_torch_available() else ()
 
     test_pruning = False
-    test_torchscript = False
     test_resize_embeddings = False
     test_head_masking = False
 
diff --git a/tests/fnet/test_modeling_fnet.py b/tests/fnet/test_modeling_fnet.py
@@ -284,7 +284,6 @@ class FNetModelTest(ModelTesterMixin, unittest.TestCase):
 
     # Skip Tests
     test_pruning = False
-    test_torchscript = False
     test_head_masking = False
     test_pruning = False
 
diff --git a/tests/glpn/test_modeling_glpn.py b/tests/glpn/test_modeling_glpn.py
@@ -150,7 +150,6 @@ class GLPNModelTest(ModelTesterMixin, unittest.TestCase):
     test_head_masking = False
     test_pruning = False
     test_resize_embeddings = False
-    test_torchscript = False
 
     def setUp(self):
         self.model_tester = GLPNModelTester(self)
diff --git a/tests/hubert/test_modeling_hubert.py b/tests/hubert/test_modeling_hubert.py
@@ -300,7 +300,6 @@ class HubertModelTest(ModelTesterMixin, unittest.TestCase):
     all_model_classes = (HubertForCTC, HubertForSequenceClassification, HubertModel) if is_torch_available() else ()
     test_pruning = False
     test_headmasking = False
-    test_torchscript = False
 
     def setUp(self):
         self.model_tester = HubertModelTester(self)
@@ -445,7 +444,6 @@ class HubertRobustModelTest(ModelTesterMixin, unittest.TestCase):
     all_model_classes = (HubertForCTC, HubertForSequenceClassification, HubertModel) if is_torch_available() else ()
     test_pruning = False
     test_headmasking = False
-    test_torchscript = False
 
     def setUp(self):
         self.model_tester = HubertModelTester(
diff --git a/tests/maskformer/test_modeling_maskformer.py b/tests/maskformer/test_modeling_maskformer.py
@@ -177,7 +177,6 @@ class MaskFormerModelTest(ModelTesterMixin, unittest.TestCase):
     all_model_classes = (MaskFormerModel, MaskFormerForInstanceSegmentation) if is_torch_available() else ()
 
     is_encoder_decoder = False
-    test_torchscript = False
     test_pruning = False
     test_head_masking = False
     test_missing_keys = False
diff --git a/tests/mpnet/test_modeling_mpnet.py b/tests/mpnet/test_modeling_mpnet.py
@@ -205,7 +205,6 @@ class MPNetModelTest(ModelTesterMixin, unittest.TestCase):
         else ()
     )
     test_pruning = False
-    test_torchscript = True
     test_resize_embeddings = True
 
     def setUp(self):
diff --git a/tests/prophetnet/test_modeling_prophetnet.py b/tests/prophetnet/test_modeling_prophetnet.py
@@ -890,7 +890,6 @@ class ProphetNetModelTest(ModelTesterMixin, GenerationTesterMixin, unittest.Test
     all_model_classes = (ProphetNetModel, ProphetNetForConditionalGeneration) if is_torch_available() else ()
     all_generative_model_classes = (ProphetNetForConditionalGeneration,) if is_torch_available() else ()
     test_pruning = False
-    test_torchscript = False
     test_resize_embeddings = False
     is_encoder_decoder = True
 
@@ -1100,7 +1099,7 @@ class ProphetNetStandaloneDecoderModelTest(ModelTesterMixin, GenerationTesterMix
     all_model_classes = (ProphetNetDecoder, ProphetNetForCausalLM) if is_torch_available() else ()
     all_generative_model_classes = (ProphetNetForCausalLM,) if is_torch_available() else ()
     test_pruning = False
-    test_torchscript = False
+
     test_resize_embeddings = False
     is_encoder_decoder = False
 
@@ -1128,7 +1127,7 @@ def test_retain_grad_hidden_states_attentions(self):
 class ProphetNetStandaloneEncoderModelTest(ModelTesterMixin, unittest.TestCase):
     all_model_classes = (ProphetNetEncoder,) if is_torch_available() else ()
     test_pruning = False
-    test_torchscript = False
+
     test_resize_embeddings = False
     is_encoder_decoder = False
 
diff --git a/tests/regnet/test_modeling_regnet.py b/tests/regnet/test_modeling_regnet.py
@@ -127,7 +127,6 @@ class RegNetModelTest(ModelTesterMixin, unittest.TestCase):
     all_model_classes = (RegNetModel, RegNetForImageClassification) if is_torch_available() else ()
 
     test_pruning = False
-    test_torchscript = False
     test_resize_embeddings = False
     test_head_masking = False
     has_attentions = False
diff --git a/tests/resnet/test_modeling_resnet.py b/tests/resnet/test_modeling_resnet.py
@@ -127,7 +127,6 @@ class ResNetModelTest(ModelTesterMixin, unittest.TestCase):
     all_model_classes = (ResNetModel, ResNetForImageClassification) if is_torch_available() else ()
 
     test_pruning = False
-    test_torchscript = False
     test_resize_embeddings = False
     test_head_masking = False
     has_attentions = False
diff --git a/tests/segformer/test_modeling_segformer.py b/tests/segformer/test_modeling_segformer.py
@@ -165,7 +165,6 @@ class SegformerModelTest(ModelTesterMixin, unittest.TestCase):
     test_head_masking = False
     test_pruning = False
     test_resize_embeddings = False
-    test_torchscript = False
 
     def setUp(self):
         self.model_tester = SegformerModelTester(self)
diff --git a/tests/sew/test_modeling_sew.py b/tests/sew/test_modeling_sew.py
@@ -303,7 +303,6 @@ class SEWModelTest(ModelTesterMixin, unittest.TestCase):
     all_model_classes = (SEWForCTC, SEWModel, SEWForSequenceClassification) if is_torch_available() else ()
     test_pruning = False
     test_headmasking = False
-    test_torchscript = False
 
     def setUp(self):
         self.model_tester = SEWModelTester(self)
diff --git a/tests/speech_to_text/test_modeling_speech_to_text.py b/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"
 
diff --git a/tests/squeezebert/test_modeling_squeezebert.py b/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
 
diff --git a/tests/swin/test_modeling_swin.py b/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
 
diff --git a/tests/t5/test_modeling_t5.py b/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 @@ def prepare_config_and_inputs_for_common(self):
 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 ()
diff --git a/tests/tapas/test_modeling_tapas.py b/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
 
diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py
@@ -617,19 +617,21 @@ def _create_and_check_torchscript(self, config, inputs_dict):
             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.")
 
diff --git a/tests/transfo_xl/test_modeling_transfo_xl.py b/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
 
diff --git a/tests/unispeech/test_modeling_unispeech.py b/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(
diff --git a/tests/van/test_modeling_van.py b/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
diff --git a/tests/vit/test_modeling_vit.py b/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
 
diff --git a/tests/wav2vec2/test_modeling_wav2vec2.py b/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(
diff --git a/tests/wavlm/test_modeling_wavlm.py b/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)

Original file line number	Diff line number	Diff line change
`@@ -192,7 +192,6 @@ class BeitModelTest(ModelTesterMixin, unittest.TestCase):`
`192`	`192`	`)`
`193`	`193`
`194`	`194`	`test_pruning = False`
`195`		`- test_torchscript = False`
`196`	`195`	`test_resize_embeddings = False`
`197`	`196`	`test_head_masking = False`
`198`	`197`
Original file line number	Diff line number	Diff line change
`@@ -219,7 +219,6 @@ class CanineModelTest(ModelTesterMixin, unittest.TestCase):`
`219`	`219`	`else ()`
`220`	`220`	`)`
`221`	`221`
`222`		`- test_torchscript = False`
`223`	`222`	`test_mismatched_shapes = False`
`224`	`223`	`test_resize_embeddings = False`
`225`	`224`	`test_pruning = False`
Original file line number	Diff line number	Diff line change
`@@ -138,7 +138,6 @@ class ConvNextModelTest(ModelTesterMixin, unittest.TestCase):`
`138`	`138`	`)`
`139`	`139`
`140`	`140`	`test_pruning = False`
`141`		`- test_torchscript = False`
`142`	`141`	`test_resize_embeddings = False`
`143`	`142`	`test_head_masking = False`
`144`	`143`	`has_attentions = False`
Original file line number	Diff line number	Diff line change
`@@ -372,7 +372,6 @@ class Data2VecAudioModelTest(ModelTesterMixin, unittest.TestCase):`
`372`	`372`	`)`
`373`	`373`	`test_pruning = False`
`374`	`374`	`test_headmasking = False`
`375`		`- test_torchscript = False`
`376`	`375`
`377`	`376`	`def setUp(self):`
`378`	`377`	`self.model_tester = Data2VecAudioModelTester(self)`
Original file line number	Diff line number	Diff line change
`@@ -169,7 +169,6 @@ class DeiTModelTest(ModelTesterMixin, unittest.TestCase):`
`169`	`169`	`)`
`170`	`170`
`171`	`171`	`test_pruning = False`
`172`		`- test_torchscript = False`
`173`	`172`	`test_resize_embeddings = False`
`174`	`173`	`test_head_masking = False`
`175`	`174`
Original file line number	Diff line number	Diff line change
`@@ -211,7 +211,6 @@ class DistilBertModelTest(ModelTesterMixin, unittest.TestCase):`
`211`	`211`	`)`
`212`	`212`	`fx_compatible = True`
`213`	`213`	`test_pruning = True`
`214`		`- test_torchscript = True`
`215`	`214`	`test_resize_embeddings = True`
`216`	`215`	`test_resize_position_embeddings = True`
`217`	`216`