def METHOD_NAME(self, x_j: Tensor, x_i: Tensor, edge_attr: Tensor,
            index: Tensor, ptr: OptTensor,
            size_i: Optional[int]) -> Tensor:
    x_j = F.leaky_relu_(self.lin1(torch.cat([x_j, edge_attr], dim=-1)))
    alpha_j = (x_j @ self.att_l.t()).squeeze(-1)
    alpha_i = (x_i @ self.att_r.t()).squeeze(-1)
    alpha = alpha_j + alpha_i
    alpha = F.leaky_relu_(alpha)
    alpha = softmax(alpha, index, ptr, size_i)
    alpha = F.dropout(alpha, p=self.dropout, training=self.training)
    return self.lin2(x_j) * alpha.unsqueeze(-1)