Skip to content

Commit 35c016c

Browse files
authored
Make it possible to pass device through pull_apart_list (#26)
1 parent b1a3594 commit 35c016c

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

merlin/core/dispatch.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -391,7 +391,7 @@ def encode_list_column(original, encoded, dtype=None):
391391
)
392392

393393

394-
def pull_apart_list(original):
394+
def pull_apart_list(original, device=None):
395395
values = flatten_list_column_values(original)
396396
if isinstance(original, pd.Series):
397397
offsets = pd.Series([0]).append(original.map(len).cumsum()).reset_index(drop=True)
@@ -402,7 +402,7 @@ def pull_apart_list(original):
402402
elements = original._column.elements
403403
if isinstance(elements, cudf.core.column.lists.ListColumn):
404404
offsets = elements.offsets[offsets]
405-
return make_series(values), make_series(offsets)
405+
return make_series(values, device), make_series(offsets, device)
406406

407407

408408
def to_arrow(x):

0 commit comments

Comments
 (0)