def extract_node_features(name, block, data, node_embed, device):
    """Extract the node features from embedding layer or raw features."""
    if name == "ogbn-mag":
        input_nodes = {
            k: v.to(device) for k, v in block.srcdata[dgl.NID].items()
        }
        # Extract node embeddings for the input nodes.
        node_features = extract_embed(node_embed, input_nodes)
        # Add the batch's raw "paper" features. Corresponds to the content
        # in the function `rel_graph_embed` comment.
        node_features.update(
            {"paper": data.node_features[("paper", "feat")].to(device)}
        )
    else:
        node_features = {
            ntype: data.node_features[(ntype, "feat")]
            for ntype in block.srctypes
        }
        # Original feature data are stored in float16 while model weights are
        # float32, so we need to convert the features to float32.
        node_features = {
            k: v.to(device).float() for k, v in node_features.items()
        }
    return node_features
Extract the node features from embedding layer or raw features.
examples/graphbolt/rgcn/hetero_rgcn.py