Only import cupy if the machine has GPUs available

smokestacklightnin · Aug 9, 2023 · db117eb · db117eb
1 parent 4be0cf6
commit db117eb
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/stardis/opacities/tests/test_voigt.py b/stardis/opacities/tests/test_voigt.py
@@ -2,7 +2,7 @@
 import numpy as np
 from math import sqrt
 from numba import cuda
-import cupy as cp
+
 from stardis.opacities.voigt import (
     faddeeva,
     _faddeeva_cuda,
@@ -13,6 +13,9 @@
 # Test cases must also take into account use of a GPU to run. If there is no GPU then the test cases will fail.
 GPUs_available = cuda.is_available()
 
+if GPUs_available:
+    import cupy as cp
+
 
 @pytest.mark.parametrize(
     "faddeeva_sample_values_input, faddeeva_sample_values_expected_result",

diff --git a/stardis/opacities/voigt.py b/stardis/opacities/voigt.py
@@ -1,9 +1,13 @@
 import numpy as np
 import numba
 from numba import cuda
-import cupy as cp
 import cmath
 
+GPUs_available = cuda.is_available()
+
+if GPUs_available:
+    import cupy as cp
+
 SQRT_PI = np.sqrt(np.pi, dtype=float)