diff --git a/distributed/protocol/cupy.py b/distributed/protocol/cupy.py index fc6e0d0871..2cc2169b07 100644 --- a/distributed/protocol/cupy.py +++ b/distributed/protocol/cupy.py @@ -85,7 +85,7 @@ def dask_deserialize_cupy_ndarray(header, frames): def reduce_matdescriptor(other): # Pickling MatDescriptor errors # xref: https://github.com/cupy/cupy/issues/3061 - return cupy.cusparse.MatDescriptor.create, () + return MatDescriptor.create, () copyreg.pickle(MatDescriptor, reduce_matdescriptor)