test_mxnet_wrapper: Feature-gate GPU test (#717)

explosion · Jul 6, 2022 · ea3c08e · ea3c08e
1 parent 7cd060e
commit ea3c08e
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/thinc/tests/layers/test_mxnet_wrapper.py b/thinc/tests/layers/test_mxnet_wrapper.py
@@ -5,7 +5,7 @@
 from thinc.api import Adam, ArgsKwargs, Model, Ops, MXNetWrapper
 from thinc.api import get_current_ops, mxnet2xp, xp2mxnet
 from thinc.types import Array2d, Array1d, IntsXd
-from thinc.compat import has_cupy, has_mxnet
+from thinc.compat import has_cupy_gpu, has_mxnet
 from thinc.util import to_categorical
 
 from ..util import check_input_converters, make_tempdir
@@ -158,7 +158,8 @@ def test_mxnet_wrapper_to_cpu(mx_model, X: Array2d):
     model.to_cpu()
 
 
-@pytest.mark.skipif(not has_mxnet or not has_cupy, reason="needs MXNet")
+@pytest.mark.skipif(not has_mxnet, reason="needs MXNet")
+@pytest.mark.skipif(not has_cupy_gpu, reason="needs GPU/cupy")
 def test_mxnet_wrapper_to_gpu(model: Model[Array2d, Array2d], X: Array2d):
     model.predict(X)
     model.to_gpu(0)