code: 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
  
imports: import argparse
  import itertools
  import sys
  import time
  import dgl
  import dgl.nn
  import numpy
  import psutil
  import torch
  import torch.nn
  import torch.nn.functional
  from dgl import AddReverse, Compose, ToSimple
  from dgl.nn import HeteroEmbedding
  from ogb.lsc import MAG240MDataset, MAG240MEvaluator
  from ogb.nodeproppred import DglNodePropPredDataset, Evaluator
  from tqdm import tqdm
label_desc: Use image node_id 4 for calling a global function with example usage: extract_node_features(name, g, input_nodes, node_embed, feats, device) and returns: node_features