[ctrl] fix pruning of MultiHeadAttention (#4904)

2020-06-10 23:36:55 +05:30
parent 4e10acb3e5
commit 5d63ca6c38
2 changed files with 21 additions and 3 deletions
--- a/tests/test_modeling_ctrl.py
+++ b/tests/test_modeling_ctrl.py
@@ -32,7 +32,7 @@ class CTRLModelTest(ModelTesterMixin, unittest.TestCase):

    all_model_classes = (CTRLModel, CTRLLMHeadModel) if is_torch_available() else ()
    all_generative_model_classes = (CTRLLMHeadModel,) if is_torch_available() else ()
-    test_pruning = False
+    test_pruning = True
    test_torchscript = False
    test_resize_embeddings = False
    test_head_masking = False