text: import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from torch_geometric.nn import SAGEConv
  
  class GNN(torch.nn.Module):
      def __init__(self, input_dim, hidden_dim, output_dim, args):
          super(GNN, self).__init__()
          self.dropout = args.dropout
          self.num_layers = args.num_layers
  
          self.convs = nn.ModuleList()
          self.bns = nn.ModuleList()
  
          self.convs.append(SAGEConv(input_dim, hidden_dim))
          self.bns.append(nn.BatchNorm1d(hidden_dim))
  
          for l in range(self.num_layers - 2):
              self.convs.append(SAGEConv(hidden_dim, hidden_dim))
              self.bns.append(nn.BatchNorm1d(hidden_dim))
          self.convs.append(SAGEConv(hidden_dim, hidden_dim))
  
          self.post_mp = nn.Linear(hidden_dim, output_dim)
  
      def forward(self, data, mode="batch"):
          if mode == "batch":
              edge_indices, x = data
              for i in range(len(self.convs) - 1):
                  edge_index = edge_indices[i]
                  x = self.convs[i](x, edge_index)
                  x = self.bns[i](x)
                  x = F.relu(x)
                  x = F.dropout(x, p=self.dropout, training=self.training)
              x = self.convs[-1](x, edge_indices[len(self.convs) - 1])
          else:
              x, edge_index = data.node_feature, data.edge_index
              for i in range(len(self.convs) - 1):
                  x = self.convs[i](x, edge_index)
                  x = self.bns[i](x)
                  x = F.relu(x)
                  x = F.dropout(x, p=self.dropout, training=self.training)
              x = self.convs[-1](x, edge_index)
          x = self.post_mp(x)
          x = F.log_softmax(x, dim=1)
          return x