Adding accelerate annotations to tests

huggingface · Apr 16, 2024 · 52695f4 · 52695f4
1 parent e55d4a8
commit 52695f4
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py
@@ -439,6 +439,8 @@ class CopyClass(model_class):
                     self.assertLessEqual(max_diff, 1e-3, msg=f"{key} not identical")
 
     @slow
+    @require_accelerate
+    @mark.accelerate_tests
     def test_save_load_low_cpu_mem_usage(self):
         config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
         with tempfile.TemporaryDirectory() as saved_model_path:
@@ -449,6 +451,8 @@ def test_save_load_low_cpu_mem_usage(self):
                 self._check_save_load_low_cpu_mem_usage(model_class, saved_model_path)
 
     @slow
+    @require_accelerate
+    @mark.accelerate_tests
     def test_save_load_low_cpu_mem_usage_checkpoints(self):
         config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
         with tempfile.TemporaryDirectory() as saved_model_path:
@@ -460,6 +464,8 @@ def test_save_load_low_cpu_mem_usage_checkpoints(self):
                 self._check_save_load_low_cpu_mem_usage(model_class, saved_model_path)
 
     @slow
+    @require_accelerate
+    @mark.accelerate_tests
     def test_save_load_low_cpu_mem_usage_no_safetensors(self):
         with tempfile.TemporaryDirectory() as saved_model_path:
             for model_class in self.all_model_classes: