Don't convert namedtuple to tuple (#1589)

* Don't convert namedtuple to tuple * Test namedtuples sent to device correctly
2026-09-24 13:50:25 +08:00 · 2020-04-30 08:04:50 -04:00
parent d40425d257
commit 3eac6cfd4f
2 changed files with 17 additions and 4 deletions
@@ -461,10 +461,15 @@ class TrainerDPMixin(ABC):

        # when tuple
        if isinstance(batch, tuple):
-            batch = list(batch)
-            for i, x in enumerate(batch):
-                batch[i] = self.__transfer_data_to_device(x, device, gpu_id)
-            return tuple(batch)
+            # when namedtuple
+            if hasattr(batch, '_fields'):
+                elem_type = type(batch)
+                return elem_type(*(self.__transfer_data_to_device(x, device, gpu_id) for x in batch))
+            else:
+                batch = list(batch)
+                for i, x in enumerate(batch):
+                    batch[i] = self.__transfer_data_to_device(x, device, gpu_id)
+                return tuple(batch)

        # when dict
        if isinstance(batch, dict):