Experiment with new mixin class

huggingface · Nov 10, 2023 · e52535c · e52535c
1 parent 400b22f
commit e52535c
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/optimum/neuron/utils/runner.py b/optimum/neuron/utils/runner.py
@@ -323,7 +323,7 @@ def check_user_logged_in_and_cache_repo_is_set(self):
         has_write_access = has_write_access_to_repo(main_repo)
         if not has_write_access:
             raise RuntimeError(
-                f"You do not have write access to {main_repo}. Please log in and/or use a custom Tranium cache repo."
+                f"You do not have write access to {main_repo}. Please log in and/or use a custom Neuron cache repo."
             )
 
     def download_model_repo_and_override_config(

diff --git a/tests/test_generate.py b/tests/test_generate.py
@@ -90,7 +90,7 @@ def _test_generative_decoding(
 class GenerateTestCase(TestCase, TrainiumTestMixin):
     @pytest.mark.skip("Remove once generate fix (#262) has been merged.")
     @is_trainium_test
-    @parameterized(GREEDY_TESTDATA)
+    @parameterized.expand(GREEDY_TESTDATA)
     def test_greedy_decoding(self, model_name, use_cache, decoder_only, compiler_flags):
         os.environ["NEURON_CC_FLAGS"] = compiler_flags
         os.environ["XLA_USE_BF16"] = "0"