def extract_node_features(name, g, input_nodes, node_embed, feats, device):
    """Extract the node features from embedding layer or raw features."""
    if name == "ogbn-mag":
        # 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": g.ndata["feat"]["paper"][input_nodes["paper"].cpu()]}
        )
        node_features = {k: e.to(device) for k, e in node_features.items()}
    else:
        node_features = {
            ntype: feats[ntype][input_nodes[ntype].cpu()].to(device)
            for ntype in input_nodes
        }
        # Original feature data are stored in float16 while model weights are
        # float32, so we need to convert the features to float32.
        # [TODO] Enable mixed precision training on GPU.
        node_features = {k: v.float() for k, v in node_features.items()}
    return node_features
Extract the node features from embedding layer or raw features.
examples/core/rgcn/hetero_rgcn.py