Skip to content

GitLab

Unverified Commit 391f513e authored by Muhammed Fatih BALIN's avatar Muhammed Fatih BALIN Committed by GitHub
Browse files

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

parent 00f33224
Show whitespace changes
Inline Side-by-side
Showing with 12 additions and 4 deletions
examples/sampling/graphbolt/link_prediction.py
View file @ 391f513e
...@@ -392,6 +392,10 @@ def main(args): ...@@ -392,6 +392,10 @@ def main(args):
dataset = gb.BuiltinDataset("ogbl-citation2").load() dataset = gb.BuiltinDataset("ogbl-citation2").load()
# Move the dataset to the selected storage. # 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) graph = dataset.graph.to(args.storage_device)
features = dataset.feature.to(args.storage_device) features = dataset.feature.to(args.storage_device)
......
examples/sampling/graphbolt/node_classification.py
View file @ 391f513e
...@@ -404,6 +404,10 @@ def main(args): ...@@ -404,6 +404,10 @@ def main(args):
dataset = gb.BuiltinDataset("ogbn-products").load() dataset = gb.BuiltinDataset("ogbn-products").load()
# Move the dataset to the selected storage. # 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) graph = dataset.graph.to(args.storage_device)
features = dataset.feature.to(args.storage_device) features = dataset.feature.to(args.storage_device)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment