def METHOD_NAME(self, nfeat, efeat, mask=None):
    """Forward computation. Note: :attr:`nfeat` and :attr:`efeat` should be
    padded with embedding of virtual nodes if :attr:`num_virtual_nodes` > 0,
    while :attr:`mask` should be padded with `0` values for virtual nodes.
    The padding should be put at the beginning.
    Parameters
    ----------
    nfeat : torch.Tensor
        A 3D input tensor. Shape: (batch_size, N, :attr:`feat_size`), where N
        is the sum of the maximum number of nodes and the number of virtual nodes.
    efeat : torch.Tensor
        Edge embedding used for attention computation and self update.
        Shape: (batch_size, N, N, :attr:`edge_feat_size`).
    mask : torch.Tensor, optional
        The attention mask used for avoiding computation on invalid
        positions, where valid positions are indicated by `0` and
        invalid positions are indicated by `-inf`.
        Shape: (batch_size, N, N). Default: None.
    Returns
    -------
    nfeat : torch.Tensor
        The output node embedding. Shape: (batch_size, N, :attr:`feat_size`).
    efeat : torch.Tensor, optional
        The output edge embedding. Shape: (batch_size, N, N, :attr:`edge_feat_size`).
        It is returned only if :attr:`edge_update` is True.
    """
    nfeat_r1 = nfeat
    efeat_r1 = efeat
    nfeat_ln = self.mha_ln_h(nfeat)
    efeat_ln = self.mha_ln_e(efeat)
    qkv = self.qkv_proj(nfeat_ln)
    e_bias = self.edge_input(efeat_ln)
    gates = self.gate(efeat_ln)
    bsz, N, _ = qkv.shape
    q_h, k_h, v_h = qkv.view(bsz, N, -1, self.num_heads).split(
        self.dot_dim, dim=2
    )
    attn_hat = torch.einsum("bldh,bmdh->blmh", q_h, k_h)
    attn_hat = attn_hat.clamp(-5, 5) + e_bias
    if mask is None:
        gates = torch.sigmoid(gates)
        attn_tild = F.softmax(attn_hat, dim=2) * gates
    else:
        gates = torch.sigmoid(gates + mask.unsqueeze(-1))
        attn_tild = F.softmax(attn_hat + mask.unsqueeze(-1), dim=2) * gates
    attn_tild = self.attn_dropout(attn_tild)
    v_attn = torch.einsum("blmh,bmkh->blkh", attn_tild, v_h)
    # Scale the aggregated values by degree.
    degrees = torch.sum(gates, dim=2, keepdim=True)
    degree_scalers = torch.log(1 + degrees)
    degree_scalers[:, : self.num_virtual_nodes] = 1.0
    v_attn = v_attn * degree_scalers
    v_attn = v_attn.reshape(bsz, N, self.num_heads * self.dot_dim)
    nfeat = self.node_output(v_attn)
    nfeat = self.mha_dropout_h(nfeat)
    nfeat.add_(nfeat_r1)
    nfeat_r2 = nfeat
    nfeat = self.node_ffn(nfeat)
    nfeat.add_(nfeat_r2)
    if self.edge_update:
        efeat = self.edge_output(attn_hat)
        efeat = self.mha_dropout_e(efeat)
        efeat.add_(efeat_r1)
        efeat_r2 = efeat
        efeat = self.edge_ffn(efeat)
        efeat.add_(efeat_r2)
        return nfeat, efeat
    return nfeat