Make it possible to pass device through pull_apart_list (#26)

NVIDIA-Merlin · Feb 28, 2022 · 35c016c · 35c016c
1 parent b1a3594
commit 35c016c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/merlin/core/dispatch.py b/merlin/core/dispatch.py
@@ -391,7 +391,7 @@ def encode_list_column(original, encoded, dtype=None):
         )
 
 
-def pull_apart_list(original):
+def pull_apart_list(original, device=None):
     values = flatten_list_column_values(original)
     if isinstance(original, pd.Series):
         offsets = pd.Series([0]).append(original.map(len).cumsum()).reset_index(drop=True)
@@ -402,7 +402,7 @@ def pull_apart_list(original):
         elements = original._column.elements
         if isinstance(elements, cudf.core.column.lists.ListColumn):
             offsets = elements.offsets[offsets]
-    return make_series(values), make_series(offsets)
+    return make_series(values, device), make_series(offsets, device)
 
 
 def to_arrow(x):