diff --git a/misinformation/test/test_multimodal_search.py b/misinformation/test/test_multimodal_search.py index 1151fad..c339e26 100644 --- a/misinformation/test/test_multimodal_search.py +++ b/misinformation/test/test_multimodal_search.py @@ -1,3 +1,4 @@ +import pytest from PIL import Image import numpy from torch import device, cuda, no_grad @@ -23,7 +24,7 @@ def test_read_img(): test_img = ms.MultimodalSearch.read_img(my_dict, TEST_IMAGE_2) assert list(numpy.array(test_img)[257][34]) == [70, 66, 63] - +@pytest.mark.skipif(not cuda.is_available(), reason="model for gpu only") def test_load_feature_extractor_model_blip2(): my_dict = {} multimodal_device = device("cuda" if cuda.is_available() else "cpu")