text: import dgl
  import torch.nn as nn
  
  """
      GAT: Graph Attention Network
      Graph Attention Networks (Veličković et al., ICLR 2018)
      https://arxiv.org/abs/1710.10903
  """
  from layers.gat_layer import GATLayer
  from layers.mlp_readout_layer import MLPReadout
  
  
  class GATNet(nn.Module):
      def __init__(self, net_params):
          super().__init__()
          num_atom_type = net_params["num_atom_type"]
          hidden_dim = net_params["hidden_dim"]
          num_heads = net_params["n_heads"]
          out_dim = net_params["out_dim"]
          in_feat_dropout = net_params["in_feat_dropout"]
          dropout = net_params["dropout"]
          n_layers = net_params["L"]
          self.readout = net_params["readout"]
          self.batch_norm = net_params["batch_norm"]
          self.residual = net_params["residual"]
  
          self.dropout = dropout
  
          self.embedding_h = nn.Embedding(num_atom_type, hidden_dim * num_heads)
          self.in_feat_dropout = nn.Dropout(in_feat_dropout)
  
          self.layers = nn.ModuleList(
              [
                  GATLayer(
                      hidden_dim * num_heads,
                      hidden_dim,
                      num_heads,
                      dropout,
                      self.batch_norm,
                      self.residual,
                  )
                  for _ in range(n_layers - 1)
              ]
          )
          self.layers.append(
              GATLayer(
                  hidden_dim * num_heads,
                  out_dim,
                  1,
                  dropout,
                  self.batch_norm,
                  self.residual,
              )
          )
          self.MLP_layer = MLPReadout(out_dim, 1)  # 1 out dim since regression problem
  
      def forward(self, g, h, e):
          h = self.embedding_h(h)
          h = self.in_feat_dropout(h)
          for conv in self.layers:
              h = conv(g, h)
          g.ndata["h"] = h
  
          if self.readout == "sum":
              hg = dgl.sum_nodes(g, "h")
          elif self.readout == "max":
              hg = dgl.max_nodes(g, "h")
          elif self.readout == "mean":
              hg = dgl.mean_nodes(g, "h")
          else:
              hg = dgl.mean_nodes(g, "h")  # default readout is mean nodes
  
          return self.MLP_layer(hg)
  
      def loss(self, scores, targets):
          loss = nn.L1Loss()(scores, targets)
          return loss
  