Fix transfer_tensors_to_cpu crash on mixed CPU/CUDA dicts

Summary: `transfer_tensors_to_cpu` selects the source device from the first CUDA tensor, but the copy loop guarded only on `isinstance(v, torch.Tensor)` and ran the device-to-host path — including `record_stream`, which is CUDA-only — on every tensor. For an input dict that mixes CPU and CUDA tensors this raised `NotImplementedError: aten::record_stream ... 'CPU' backend`. Guard the copy loop so only CUDA tensors are transferred; CPU tensors and non-tensor values pass through unchanged, matching the existing source-device selection guard.

Differential Revision: D113422068

Author: jeffkbkim

torchrec/metrics/deferrable_metrics.py

@@ -77,8 +77,9 @@ def transfer_tensors_to_cpu(
 
     Returns the CPU tensor dict and a CUDA event recorded on the dedicated
     stream that tracks completion. For CPU-only inputs, returns the dict
-    as-is with None event. Non-tensor values are preserved unchanged.
-    All tensors are assumed to be on the same device.
+    as-is with None event. Only CUDA tensors are transferred; CPU tensors and
+    non-tensor values are passed through unchanged. CUDA tensors are assumed to
+    share a single device.
     """
     source_device: torch.device | None = None
     for v in tensors.values():
@@ -99,7 +100,7 @@ def transfer_tensors_to_cpu(
         with torch.cuda.stream(dtoh_stream):
             cpu_tensors: dict[str, Any] = {}
             for k, v in tensors.items():
-                if isinstance(v, torch.Tensor):
+                if isinstance(v, torch.Tensor) and v.device.type == "cuda":
                     dst = torch.empty(
                         v.shape, dtype=v.dtype, device="cpu", pin_memory=True
                     )

torchrec/metrics/tests/test_deferrable_metrics.py

@@ -396,6 +396,31 @@ def test_cuda_transfer_uses_dedicated_stream(self) -> None:
         dtoh_stream.synchronize()
         self.assertTrue(event.query())
 
+    @unittest.skipIf(not torch.cuda.is_available(), "CUDA required")
+    def test_mixed_cpu_and_cuda_tensors(self) -> None:
+        """Regression: a dict mixing CPU and CUDA tensors must not crash.
+        Previously the copy loop guarded only on isinstance(v, Tensor) and
+        called v.record_stream (CUDA-only) on the CPU tensors too, raising
+        NotImplementedError. CPU tensors must pass through untouched while the
+        CUDA tensor is transferred to CPU."""
+        cpu_tensor = torch.tensor([1.0, 2.0])
+        tensors: dict[str, Any] = {
+            "cuda_val": torch.tensor([0.5], device="cuda"),
+            "cpu_val": cpu_tensor,
+            "name": "task1",
+        }
+        cpu_tensors, event = transfer_tensors_to_cpu(tensors)
+        self.assertIsNotNone(event)
+        # CUDA tensor transferred to CPU.
+        self.assertEqual(cpu_tensors["cuda_val"].device.type, "cpu")
+        # CPU tensor passed through as the same object (not copied/pinned).
+        self.assertIs(cpu_tensors["cpu_val"], cpu_tensor)
+        # Non-tensor preserved.
+        self.assertEqual(cpu_tensors["name"], "task1")
+        if event is not None:
+            event.synchronize()
+        torch.testing.assert_close(cpu_tensors["cuda_val"], torch.tensor([0.5]))
+
     @unittest.skipIf(not torch.cuda.is_available(), "CUDA required")
     def test_cuda_transfer_destination_is_pinned(self) -> None:
         """The CPU destination must be pinned so cudaMemcpyAsync is truly