Skip to content

GitLab

Unverified Commit dfa32286 authored by Ramon Zhou's avatar Ramon Zhou Committed by GitHub
Browse files

[Misc] DGL multi-gpu example run error work around (#6999)

parent d2497448
Hide whitespace changes
Inline Side-by-side
Showing with 10 additions and 6 deletions
examples/multigpu/node_classification_sage.py
View file @ dfa32286
...@@ -171,12 +171,16 @@ def train( ...@@ -171,12 +171,16 @@ def train(
use_uva, use_uva,
): ):
# Instantiate a neighbor sampler # Instantiate a neighbor sampler
sampler = NeighborSampler( if args.mode == "benchmark":
[10, 10, 10], # A work-around to prevent CUDA running error. For more details, please
prefetch_node_feats=["feat"], # see https://github.com/dmlc/dgl/issues/6697.
prefetch_labels=["label"], sampler = NeighborSampler([10, 10, 10], fused=False)
fused=(args.mode != "benchmark"), else:
) sampler = NeighborSampler(
[10, 10, 10],
prefetch_node_feats=["feat"],
prefetch_labels=["label"],
)
train_dataloader = DataLoader( train_dataloader = DataLoader(
g, g,
train_idx, train_idx,
......
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