def forward(self, h_src, h_dst):
        r"""

        Description
        -----------
        Predict for pairs of node representations.

        Parameters
        ----------
        h_src : torch.Tensor
            Source node features. The tensor is of shape :math:`(E, D_{in})`,
            where :math:`E` is the number of edges/node pairs, and :math:`D_{in}`
            is the input feature size.
        h_dst : torch.Tensor
            Destination node features. The tensor is of shape :math:`(E, D_{in})`,
            where :math:`E` is the number of edges/node pairs, and :math:`D_{in}`
            is the input feature size.

        Returns
        -------
        torch.Tensor
            The output features.
        """
        if self.op == "dot":
            N, D = h_src.shape
            h = torch.bmm(h_src.view(N, 1, D), h_dst.view(N, D, 1)).squeeze(-1)
        elif self.op == "cos":
            h = F.cosine_similarity(h_src, h_dst).unsqueeze(-1)
        elif self.op == "ele":
            h = h_src * h_dst
        else:
            h = torch.cat([h_src, h_dst], dim=-1)

        if self.linear is not None:
            h = self.linear(h)

        return h


        Description
        -----------
        Predict for pairs of node representations.

        Parameters
        ----------
        h_src : torch.Tensor
            Source node features. The tensor is of shape :math:`(E, D_{in})`,
            where :math:`E` is the number of edges/node pairs, and :math:`D_{in}`
            is the input feature size.
        h_dst : torch.Tensor
            Destination node features. The tensor is of shape :math:`(E, D_{in})`,
            where :math:`E` is the number of edges/node pairs, and :math:`D_{in}`
            is the input feature size.

        Returns
        -------
        torch.Tensor
            The output features.
        
python/dgl/nn/pytorch/link/edgepred.py