Skip to content

Commit

Permalink
check cuda in class decorator
Browse files Browse the repository at this point in the history
  • Loading branch information
AlexanderDokuchaev committed Oct 12, 2023
1 parent b1536d2 commit 5bafdab
Show file tree
Hide file tree
Showing 2 changed files with 2 additions and 6 deletions.
5 changes: 1 addition & 4 deletions tests/torch/ptq/test_fast_bias_correction.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 1 addition & 2 deletions tests/torch/test_tensor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down

0 comments on commit 5bafdab

Please sign in to comment.