ci: mark model_parallel tests as cuda specific

`parallelize()` API is deprecated in favor of accelerate's `device_map="auto"` and therefore is not accepting new features. At the same time `parallelize()` implementation is currently CUDA-specific. This commit marks respective ci tests with `@require_torch_gpu`. Fixes: #35252 Signed-off-by: Dmitry Rogozhkin <[email protected]>
huggingface · Dec 13, 2024 · 4aadd0b · 4aadd0b
1 parent add53e2
commit 4aadd0b
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py
@@ -3046,6 +3046,7 @@ def test_multi_gpu_data_parallel_forward(self):
             with torch.no_grad():
                 _ = model(**self._prepare_for_class(inputs_dict, model_class))
 
+    @require_torch_gpu
     @require_torch_multi_gpu
     def test_model_parallelization(self):
         if not self.test_model_parallel:
@@ -3108,6 +3109,7 @@ def get_current_gpu_memory_use():
             gc.collect()
             torch.cuda.empty_cache()
 
+    @require_torch_gpu
     @require_torch_multi_gpu
     def test_model_parallel_equal_results(self):
         if not self.test_model_parallel: