diff --git a/tests/kernelcreation/test_reduction.py b/tests/kernelcreation/test_reduction.py index 992c328d74310090f170b580098dc9971e5f0e98..c3775964bccf556cb7a28d16b673168a92c3ecd2 100644 --- a/tests/kernelcreation/test_reduction.py +++ b/tests/kernelcreation/test_reduction.py @@ -2,7 +2,6 @@ import pytest import numpy as np import pystencils as ps -from cupy_backends.cuda.api.runtime import CUDARuntimeError from pystencils.sympyextensions import reduction_assignment_from_str INIT_W = 5 @@ -51,6 +50,7 @@ def test_reduction_cpu(instruction_set, dtype, op): def test_reduction_gpu(dtype, op): try: import cupy as cp + from cupy_backends.cuda.api.runtime import CUDARuntimeError device_count = range(cp.cuda.runtime.getDeviceCount()) print(f"Found {device_count} GPUs")