def create_dataloader(
    graph,
    features,
    itemset,
    batch_size,
    fanout,
    device,
    job,
):
    """Create a GraphBolt dataloader for training, validation or testing."""
    datapipe = gb.ItemSampler(
        itemset,
        batch_size=batch_size,
        shuffle=(job == "train"),
        drop_last=(job == "train"),
    )
    need_copy = True
    # Copy the data to the specified device.
    if args.graph_device != "cpu" and need_copy:
        datapipe = datapipe.copy_to(device=device)
        need_copy = False
    # Sample neighbors for each node in the mini-batch.
    datapipe = getattr(datapipe, args.sample_mode)(
        graph,
        fanout if job != "infer" else [-1],
        overlap_fetch=args.overlap_graph_fetch,
        num_gpu_cached_edges=args.num_gpu_cached_edges,
        gpu_cache_threshold=args.gpu_graph_caching_threshold,
        asynchronous=args.graph_device != "cpu",
    )
    # Copy the data to the specified device.
    if args.feature_device != "cpu" and need_copy:
        datapipe = datapipe.copy_to(device=device)
        need_copy = False

    node_feature_keys = {"paper": ["feat"], "author": ["feat"]}
    if args.dataset == "ogb-lsc-mag240m":
        node_feature_keys["institution"] = ["feat"]
    if "igb-het" in args.dataset:
        node_feature_keys["institute"] = ["feat"]
        node_feature_keys["fos"] = ["feat"]
    # Fetch node features for the sampled subgraph.
    datapipe = datapipe.fetch_feature(
        features,
        node_feature_keys,
        overlap_fetch=args.overlap_feature_fetch,
    )

    # Copy the data to the specified device.
    if need_copy:
        datapipe = datapipe.copy_to(device=device)
    # Create and return a DataLoader to handle data loading.
    return gb.DataLoader(datapipe, num_workers=args.num_workers)
Create a GraphBolt dataloader for training, validation or testing.
examples/graphbolt/pyg/hetero/node_classification.py