7be56405f2af0f700763249da3c368976582d792,examples/tgn.py,GraphAttentionEmbedding,forward,#GraphAttentionEmbedding#,55

Before Change

        self.time_enc = time_enc

    def forward(self, x, adj_t, t):
        if adj_t.nnz() > 0:
            _, _, value = adj_t.coo()
            
            edge_t = value[:, 0].squeeze()
            rel_t = edge_t - t
            rel_t_enc = self.time_enc(rel_t.float())
            
            edge_feat = value[:, 1:]
            
            edge_attr = torch.cat([rel_t_enc, edge_feat], dim=1)
            adj_t.set_value_(edge_attr, layout="coo")
            x = self.conv((x, x[:adj_t.size(0)]), adj_t)
        else:
            x = self.conv.lin_skip(x) 
        return x


class LinkPredictor(torch.nn.Module):
    def __init__(self, in_channels):

After Change

    def forward(self, x, last_update, edge_index, t, msg):
        rel_t = last_update[edge_index[0]] - t
        rel_t_enc = self.time_enc(rel_t.to(x.dtype))
        edge_attr = torch.cat([rel_t_enc, msg], dim=-1)
        return self.conv(x, edge_index, edge_attr)


class LinkPredictor(torch.nn.Module):
    def __init__(self, in_channels):

Instances