Fix #146 (#149)

* Fix test * Add init file * Style
huggingface · Jul 24, 2023 · c102a95 · c102a95
1 parent ae2a7b5
commit c102a95
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 1 deletion.
diff --git a/optimum/exporters/neuron/__main__.py b/optimum/exporters/neuron/__main__.py
@@ -34,7 +34,7 @@
     is_neuronx_available,
 )
 from ...neuron.utils.version_utils import check_compiler_compatibility_for_stable_diffusion
-from ...utils import logging, is_diffusers_available
+from ...utils import is_diffusers_available, logging
 from ...utils.save_utils import maybe_save_preprocessors
 from ..error_utils import AtolError, OutputMatchError, ShapeError
 from ..tasks import TasksManager

diff --git a/tests/distributed/__init__.py b/tests/distributed/__init__.py
diff --git a/tests/distributed/test_utils.py b/tests/distributed/test_utils.py
@@ -31,6 +31,8 @@
 )
 from optimum.neuron.utils.patching import patch_everywhere
 
+from ..test_utils import is_trainium_test
+
 
 def test_load_tensor_for_weight():
     with TemporaryDirectory() as tmpdirname:
@@ -59,6 +61,7 @@ def test_load_tensor_for_weight():
         assert loaded_sliced_t1.numel() == loaded_sliced_t1.storage().size()
 
 
+@is_trainium_test
 class ParallelUtilsTestCase(unittest.TestCase):
     TP_GROUP = 0
     TP_SIZE = 8