Fix pin_memory_fn for NamedTuple (pytorch#1086)

Summary:
Fixes pytorch#1085

Per title. And, even though I can add a test, this test won't be executed as we don't have a GPU CI machine yet.

I have tested on my local machine though

Pull Request resolved: pytorch#1086

Reviewed By: NivekT

Differential Revision: D44094225

Pulled By: ejguan

fbshipit-source-id: 9c8414c31b76c93cee7e31c4e2da14076e9792bf

test/test_iterdatapipe.py

@@ -13,7 +13,7 @@
 
 from collections import defaultdict
 from functools import partial
-from typing import Dict
+from typing import Dict, NamedTuple
 
 import expecttest
 import torch
@@ -91,6 +91,11 @@ async def _async_x_mul_y(x, y):
     return x * y
 
 
+class NamedTensors(NamedTuple):
+    x: torch.Tensor
+    y: torch.Tensor
+
+
 class TestIterDataPipe(expecttest.TestCase):
     def test_in_memory_cache_holder_iterdatapipe(self) -> None:
         source_dp = IterableWrapper(range(10))
@@ -1521,6 +1526,10 @@ def test_pin_memory(self):
         dp = IterableWrapper([{str(i): (i, i + 1)} for i in range(10)]).map(_convert_to_tensor).pin_memory()
         self.assertTrue(all(v.is_pinned() for d in dp for v in d.values()))
 
+        # NamedTuple
+        dp = IterableWrapper([NamedTensors(torch.tensor(i), torch.tensor(i + 1)) for i in range(10)]).pin_memory()
+        self.assertTrue(all(v.is_pinned() for d in dp for v in d))
+
         # Dict of List of Tensors
         dp = (
             IterableWrapper([{str(i): [(i - 1, i), (i, i + 1)]} for i in range(10)])

torchdata/datapipes/utils/pin_memory.py

@@ -27,7 +27,7 @@ def pin_memory_fn(data, device=None):
     elif isinstance(data, collections.abc.Sequence):
         pinned_data = [pin_memory_fn(sample, device) for sample in data]  # type: ignore[assignment]
         try:
-            type(data)(*pinned_data)
+            return type(data)(*pinned_data)
         except TypeError:
             # The sequence type may not support `__init__(iterable)` (e.g., `range`).
             return pinned_data