class GraphAttentionLayer(nn.Module):
    def __init__(self, emb_dim=256, ff_dim=1024):
        super(GraphAttentionLayer, self).__init__()
        self.linear1 = nn.Linear(emb_dim, ff_dim)
        self.linear2 = nn.Linear(ff_dim, emb_dim)
        self.gat = GATConv(emb_dim, emb_dim, heads=8, dropout=0.2, add_self_loops=False)
        
    def forward(self, x, edge_index, edge_attr):
        # x.shape >>> num_node, feat_dim
        # edge_index.shape >>> 2, num_edges
        # edge_attr.shape >>> num_edges, feat_dim2
        x = self.gat(x, edge_index, edge_attr)
        
        return x