Update

luxonis · Jan 12, 2024 · b937ac8 · b937ac8
1 parent 52dd127
commit b937ac8
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git a/tests/integration/test_pipeline.py b/tests/integration/test_pipeline.py
@@ -378,7 +378,6 @@ def test_cuda_simple_sdxl_turbo_classification_pipeline():
     _check_detection_pipeline(cmd, target_folder)
 
 
-
 # def test_cpu_simple_sdxl_classification_pipeline():
 #     # Define target folder
 #     target_folder = "data/data-cls-cpu-simple-sdxl/"

diff --git a/tests/unittests/test_image_generation.py b/tests/unittests/test_image_generation.py
@@ -50,13 +50,13 @@ def _check_image_generator(image_generator_class: Type[Union[StableDiffusionImag
     image_generator.release(empty_cuda_cache=True if device != "cpu" else False)
 
 
-# @pytest.mark.skipif(not torch.cuda.is_available(), reason="Test requires GPU")
-# def test_cuda_sdxl_image_generator():
-#     _check_image_generator(StableDiffusionImageGenerator, "cuda")
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Test requires GPU")
+def test_cuda_sdxl_image_generator():
+    _check_image_generator(StableDiffusionImageGenerator, "cuda")
 
 
-# def test_cpu_sdxl_image_generator():
-#     _check_image_generator(StableDiffusionImageGenerator, "cpu")
+def test_cpu_sdxl_image_generator():
+    _check_image_generator(StableDiffusionImageGenerator, "cpu")
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="Test requires GPU")