diff --git a/tests/sparseml/test_clear_ml.py b/tests/sparseml/transformers/test_clear_ml.py similarity index 78% rename from tests/sparseml/test_clear_ml.py rename to tests/sparseml/transformers/test_clear_ml.py index 987d15a15fe..24ec38e2b51 100644 --- a/tests/sparseml/test_clear_ml.py +++ b/tests/sparseml/transformers/test_clear_ml.py @@ -16,27 +16,11 @@ from clearml import Task from sparseml.transformers import apply -from sparseml.utils import is_package_available - - -is_torch_available = is_package_available("torch") -if is_torch_available: - import torch - - torch_err = None -else: - torch = object - torch_err = ModuleNotFoundError( - "`torch` is not installed, use `pip install torch` to log to Weights and Biases" - ) def test_oneshot_and_finetune(tmp_path: Path): recipe_str = "tests/sparseml/transformers/finetune/test_alternate_recipe.yaml" model = "Xenova/llama2.c-stories15M" - device = "cuda:0" - if is_torch_available and not torch.cuda.is_available(): - device = "cpu" dataset = "wikitext" dataset_config_name = "wikitext-2-raw-v1" concatenate_data = True @@ -59,5 +43,4 @@ def test_oneshot_and_finetune(tmp_path: Path): max_steps=max_steps, concatenate_data=concatenate_data, splits=splits, - oneshot_device=device, )