Update python/dask_cudf/dask_cudf/backends.py

Co-authored-by: Mads R. B. Kristensen <[email protected]>
rapidsai · Nov 11, 2021 · e5d726b · e5d726b
1 parent 4373a95
commit e5d726b
Showing 1 changed file with 8 additions and 9 deletions.
diff --git a/python/dask_cudf/dask_cudf/backends.py b/python/dask_cudf/dask_cudf/backends.py
@@ -348,12 +348,11 @@ def group_split_cudf(df, c, k, ignore_index=False):
     )
 
 
-@sizeof_dispatch.register_lazy("cudf")
-def register_cudf():
-    @sizeof_dispatch.register((cudf.DataFrame))
-    def sizeof_cudf_dataframe(df):
-        return int(df.memory_usage().sum())
-
-    @sizeof_dispatch.register((cudf.Series, cudf.BaseIndex))
-    def sizeof_cudf_series_index(obj):
-        return obj.memory_usage()
+@sizeof_dispatch.register(cudf.DataFrame)
+def sizeof_cudf_dataframe(df):
+    return int(df.memory_usage().sum())
+
+
+@sizeof_dispatch.register((cudf.Series, cudf.BaseIndex))
+def sizeof_cudf_series_index(obj):
+    return obj.memory_usage()