fix tests

huggingface · Oct 22, 2024 · 38cda01 · 38cda01
1 parent 8955a24
commit 38cda01
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/optimum/exporters/neuron/__main__.py b/optimum/exporters/neuron/__main__.py
@@ -261,10 +261,10 @@ def infer_stable_diffusion_shapes_from_diffusers(
 def get_submodels_and_neuron_configs(
     model: Union["PreTrainedModel", "DiffusionPipeline"],
     input_shapes: Dict[str, int],
-    tensor_parallel_size: int,
     task: str,
     output: Path,
     library_name: str,
+    tensor_parallel_size: int = 1,
     subfolder: str = "",
     dynamic_batch_size: bool = False,
     model_name_or_path: Optional[Union[str, Path]] = None,

diff --git a/tests/cli/test_export_cli.py b/tests/cli/test_export_cli.py
@@ -303,6 +303,9 @@ def test_replace_unet(self):
                 check=True,
             )
 
+    @unittest.skip(
+        "T5 compilation broken since neuron sdk 2.20, wait for the fix: https://github.com/aws-neuron/aws-neuron-sdk/issues/1013."
+    )
     @requires_neuronx
     def test_encoder_decoder(self):
         model_id = "hf-internal-testing/tiny-random-t5"
@@ -332,6 +335,9 @@ def test_encoder_decoder(self):
                 check=True,
             )
 
+    @unittest.skip(
+        "T5 compilation broken since neuron sdk 2.20, wait for the fix: https://github.com/aws-neuron/aws-neuron-sdk/issues/1013."
+    )
     @requires_neuronx
     def test_encoder_decoder_optional_outputs(self):
         model_id = "hf-internal-testing/tiny-random-t5"