Return container itself from singleton move_to.

2026-06-27 17:14:53 +08:00 · 2022-08-27 05:03:12 -04:00
parent 5390101e06
commit 0df3359918
1 changed files with 3 additions and 1 deletions
@@ -50,11 +50,13 @@ def move_to(device, *containers):
                g = enumerate(container)
            for i, v in g:
                moved = move_to(device, v)
-                if moved is not None:
+                if moved is not None and moved is not v:
                    container[i] = moved
        elif isinstance(container, (torch.nn.Parameter, torch.Tensor)):
            assert len(containers) == 1, 'Use move_to to move tensors in containers.'
            return container.to(device)
+    if len(containers) == 1:
+        return container


 def length_collation(token_size):