diff --git a/thinc/backends/cupy_ops.py b/thinc/backends/cupy_ops.py index 1e1e5b92b..472b6c542 100644 --- a/thinc/backends/cupy_ops.py +++ b/thinc/backends/cupy_ops.py @@ -1,7 +1,7 @@ import numpy from .. import registry -from ..compat import cupy, cupyx +from ..compat import cublas, cupy, cupyx from ..types import DeviceTypes from ..util import ( is_cupy_array, @@ -257,7 +257,7 @@ def clip_gradient(self, gradient, threshold): # implementation. def frobenius_norm(X): X_vec = X.reshape(-1) - return cupy.cublas.nrm2(X_vec) + return cublas.nrm2(X_vec) grad_norm = cupy.maximum(frobenius_norm(gradient), 1e-12) gradient *= cupy.minimum(threshold, grad_norm) / grad_norm diff --git a/thinc/compat.py b/thinc/compat.py index 5d600796a..c7b47cbe6 100644 --- a/thinc/compat.py +++ b/thinc/compat.py @@ -4,9 +4,11 @@ try: # pragma: no cover import cupy + import cupy.cublas import cupyx has_cupy = True + cublas = cupy.cublas cupy_version = Version(cupy.__version__) try: cupy.cuda.runtime.getDeviceCount() @@ -20,6 +22,7 @@ else: cupy_from_dlpack = cupy.fromDlpack except (ImportError, AttributeError): + cublas = None cupy = None cupyx = None cupy_version = Version("0.0.0")