[GraphBolt][CUDA] Make sure in-place pinning is used in examples. (#7138)

examples/sampling/graphbolt/link_prediction.py

@@ -392,6 +392,10 @@ def main(args):
     dataset = gb.BuiltinDataset("ogbl-citation2").load()
 
     # Move the dataset to the selected storage.
+    if args.storage_device == "pinned":
+        graph = dataset.graph.pin_memory_()
+        features = dataset.feature.pin_memory_()
+    else:
         graph = dataset.graph.to(args.storage_device)
         features = dataset.feature.to(args.storage_device)
 

examples/sampling/graphbolt/node_classification.py

@@ -404,6 +404,10 @@ def main(args):
     dataset = gb.BuiltinDataset("ogbn-products").load()
 
     # Move the dataset to the selected storage.
+    if args.storage_device == "pinned":
+        graph = dataset.graph.pin_memory_()
+        features = dataset.feature.pin_memory_()
+    else:
         graph = dataset.graph.to(args.storage_device)
         features = dataset.feature.to(args.storage_device)