move the is_pinned check insided CopyTo.

python/dgl/graphbolt/base.py

@@ -319,8 +319,6 @@ def apply_to(x, device, non_blocking=False):
         return x
     if not non_blocking:
         return x.to(device)
-    # The copy is non blocking only if the objects are pinned.
-    assert x.is_pinned(), f"{x} should be pinned."
     return x.to(device, non_blocking=True)
 
 
@@ -373,6 +371,9 @@ def __init__(self, datapipe, device, non_blocking=False):
 
     def __iter__(self):
         for data in self.datapipe:
+            if self.non_blocking:
+                # The copy is non blocking only if contents of data are pinned.
+                assert data.is_pinned(), f"{data} should be pinned."
             yield recursive_apply(
                 data, apply_to, self.device, self.non_blocking
             )