code: def _logits(self, graph_mb_data):
      graph_output: GnnOutput = self.__gnn(**graph_mb_data)
      # Gather the output representation of the nodes of interest
      supernode_idxs = graph_output.reference_nodes_idx["supernodes"]
      supernode_graph_idx = graph_output.reference_nodes_graph_idx[
          "supernodes"
      ]
      supernode_representations = (
          graph_output.output_node_representations[supernode_idxs]
      )  # [num_supernodes_in_mb, D]
      return (
          self.__node_to_class(supernode_representations),
          supernode_graph_idx,
      )
  
imports: from typing_extensions import TypedDict
  import numpy
  import torch
  from dpu_utils.mlutils import Vocabulary
  from torch import nn
  from typing import Any, Counter, Dict, Iterator, List, NamedTuple, Optional, Tuple, Union
  from ptgnn.baseneuralmodel import AbstractNeuralModel, ModuleWithMetrics
  from ptgnn.baseneuralmodel.utils.data import enforce_not_None
  from ptgnn.neuralmodels.gnn import GnnOutput, GraphData, TensorizedGraphData
  from ptgnn.neuralmodels.gnn.graphneuralnetwork import GraphNeuralNetwork, GraphNeuralNetworkModel
label_desc: Use image node_id 4 for calling the Graph2ClassModule obj's underlying member method code with example usage: obj._logits(graph_mb_data) and returns: self,  supernode_graph_idx