code: def get_batch_id(num_nodes: torch.Tensor):
      """Convert the num_nodes array obtained from batch graph to batch_id array
      for each node.
  
      Args:
          num_nodes (torch.Tensor): The tensor whose element is the number of nodes
              in each graph in the batch graph.
      """
      batch_size = num_nodes.size(0)
      batch_ids = []
      for i in range(batch_size):
          item = torch.full(
              (num_nodes[i],),
              i,
              dtype=torch.long,
              device=num_nodes.device,
          )
          batch_ids.append(item)
      return torch.cat(batch_ids)
  
imports: import logging
  import math
  import torch
  from scipy.stats import t
label_desc: Use image node_id 2 for calling a global function with example usage: get_batch_id(num_nodes) and returns: torch