Some questions about egnn

[Byun-jinyoung]

Hi,

I wonder about the implementation of two parts in the egnn codes.

First, in E_GCL code,

    def forward(self, h, edge_index, coord, edge_attr=None, node_attr=None):
        row, col = edge_index
        radial, coord_diff = self.coord2radial(edge_index, coord)

        edge_feat = self.edge_model(h[row], h[col], radial, edge_attr)
        coord     = self.coord_model(coord, edge_index, coord_diff, edge_feat)
        h, agg    = self.node_model(h, edge_index, edge_feat, node_attr)
        # coord = self.node_coord_model(h, coord)                                                                                
        # x = self.node_model(x, edge_index, x[col], u, batch)  # GCN                                                            
        return h, coord, edge_attr

the edge_feat is not return of forward method and the edge_attr, which is input of forward, is just returned.
Can I use the edge_feat as the one of the returns of forward method?

Second, Can I predict the h, edge_attr with the EGNN Autoencoder?
In example main_ae.py code, the decoder simply consists of linear transformation and just predicts adjacency matrix.
What I want to do is to predict the h_x (node attribute) and edge attribute. So did I just use the output of autoencoder as same way of adjacency matrix?

Best regards,
Jinyoung