4511a4ef8da6432f137f2ebdbb2ab78a617215f3,torch_geometric/nn/conv/nn_conv.py,NNConv,forward,#NNConv#Any#Any#Any#,77

Before Change

        
        x = x.unsqueeze(-1) if x.dim() == 1 else x
        pseudo = edge_attr.unsqueeze(-1) if edge_attr.dim() == 1 else edge_attr
        return self.propagate(edge_index, x=x, pseudo=pseudo)

    def message(self, x_j, pseudo):
        weight = self.nn(pseudo).view(-1, self.in_channels, self.out_channels)
        return torch.matmul(x_j.unsqueeze(1), weight).squeeze(1)

After Change

        
        x = x.unsqueeze(-1) if x.dim() == 1 else x
        pseudo = edge_attr.unsqueeze(-1) if edge_attr.dim() == 1 else edge_attr
        out = self.propagate(edge_index, x=x, pseudo=pseudo)
        if self.root is not None:
            out += torch.matmul(x, self.root)
        if self.bias is not None:
            out += self.bias
        return out

    def message(self, x_j, pseudo):
        weight = self.nn(pseudo).view(-1, self.in_channels, self.out_channels)
        return torch.matmul(x_j.unsqueeze(1), weight).squeeze(1)

Instances