Skip to content

Commit bc08a37

Browse files
committed
fix(nyz): fix to_device bug when using ttorch (#561)
1 parent c50379c commit bc08a37

File tree

1 file changed

+2
-0
lines changed

1 file changed

+2
-0
lines changed

ding/torch_utils/data_helper.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@ def to_device(item: Any, device: str, ignore_keys: list = []) -> Any:
3030
"""
3131
if isinstance(item, torch.nn.Module):
3232
return item.to(device)
33+
elif isinstance(item, ttorch.Tensor):
34+
return item.to(device)
3335
elif isinstance(item, torch.Tensor):
3436
return item.to(device)
3537
elif isinstance(item, Sequence):

0 commit comments

Comments
 (0)