From 5bafdab392a37190279a2d5f205ed5c05410dcaf Mon Sep 17 00:00:00 2001 From: Alexander Dokuchaev Date: Thu, 12 Oct 2023 18:31:34 +0300 Subject: [PATCH] check cuda in class decorator --- tests/torch/ptq/test_fast_bias_correction.py | 5 +---- tests/torch/test_tensor.py | 3 +-- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/torch/ptq/test_fast_bias_correction.py b/tests/torch/ptq/test_fast_bias_correction.py index 7d4779044a1..e5d3b937376 100644 --- a/tests/torch/ptq/test_fast_bias_correction.py +++ b/tests/torch/ptq/test_fast_bias_correction.py @@ -62,17 +62,14 @@ def check_bias(model: NNCFNetwork, ref_bias: list): raise ValueError("Not found node with bias") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Skipping for CPU-only setups") class TestTorchCudaFBCAlgorithm(TestTorchFBCAlgorithm): @staticmethod def list_to_backend_type(data: List) -> torch.Tensor: - if not torch.cuda.is_available(): - pytest.skip("Skipping for CPU-only setups") return torch.Tensor(data).cuda() @staticmethod def backend_specific_model(model: bool, tmp_dir: str): - if not torch.cuda.is_available(): - pytest.skip("Skipping for CPU-only setups") return get_nncf_network(model.cuda(), model.INPUT_SIZE) @staticmethod diff --git a/tests/torch/test_tensor.py b/tests/torch/test_tensor.py index eb4d907022b..35402c44891 100644 --- a/tests/torch/test_tensor.py +++ b/tests/torch/test_tensor.py @@ -22,11 +22,10 @@ def to_tensor(x): return torch.tensor(x) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Skipping for CPU-only setups") class TestCudaPTNNCFTensorOperators(TemplateTestNNCFTensorOperators): @staticmethod def to_tensor(x): - if not torch.cuda.is_available(): - pytest.skip("Skipping for CPU-only setups") return torch.tensor(x).cuda() def test_device(self):