sizeof.register_lazy("cupy"): Includes cupy.ndarray handle

rapidsai · Dec 2, 2020 · 159fa02 · 159fa02
1 parent 4671da6
commit 159fa02
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/dask_cuda/get_device_memory_objects.py b/dask_cuda/get_device_memory_objects.py
@@ -75,9 +75,13 @@ def get_device_memory_objects_cudf_index(obj):
 
 
 @sizeof.register_lazy("cupy")
-def register_cupy():
-    from cupy.cuda.memory import BaseMemory
+def register_cupy():  # NB: this overwrites dask.sizeof.register_cupy()
+    import cupy.cuda.memory
 
-    @sizeof.register(BaseMemory)
-    def sizeof_cupy_ndarray(x):
+    @sizeof.register(cupy.cuda.memory.BaseMemory)
+    def sizeof_cupy_base_memory(x):
         return int(x.size)
+
+    @sizeof.register(cupy.ndarray)
+    def sizeof_cupy_ndarray(x):
+        return int(x.nbytes)