def METHOD_NAME(
    self,
    input: Union[torch.Tensor, ComplexTensor],
    ilens: torch.Tensor,
    bottleneck_feat: torch.Tensor,
    num_spk: int,
) -> Tuple[List[Union[torch.Tensor, ComplexTensor]], torch.Tensor, OrderedDict]:
    """Keep this API same with TasNet.
    Args:
        input: [M, K, N], M is batch size
        ilens (torch.Tensor): (M,)
        bottleneck_feat: [M, K, B]
        num_spk: number of speakers
        (Training: oracle,
        Inference: estimated by other module (e.g, EEND-EDA))
    Returns:
        masked (List[Union(torch.Tensor, ComplexTensor)]): [(M, K, N), ...]
        ilens (torch.Tensor): (M,)
        others predicted data, e.g. masks: OrderedDict[
            'mask_spk1': torch.Tensor(Batch, Frames, Freq),
            'mask_spk2': torch.Tensor(Batch, Frames, Freq),
            ...
            'mask_spkn': torch.Tensor(Batch, Frames, Freq),
        ]
    """
    M, K, N = input.size()
    bottleneck_feat = bottleneck_feat.transpose(1, 2)  # [M, B, K]
    score = self.mask_conv1x1[num_spk - 1](
        bottleneck_feat
    )  # [M, B, K] -> [M, num_spk*N, K]
    # add other outputs of the module list with factor 0.0
    # to enable distributed training
    for z in range(self._max_num_spk):
        if z != num_spk - 1:
            score += 0.0 * F.interpolate(
                self.mask_conv1x1[z](bottleneck_feat).transpose(1, 2),
                size=num_spk * N,
            ).transpose(1, 2)
    score = score.view(M, num_spk, N, K)  # [M, num_spk*N, K] -> [M, num_spk, N, K]
    if self.mask_nonlinear == "softmax":
        est_mask = F.softmax(score, dim=1)
    elif self.mask_nonlinear == "relu":
        est_mask = F.relu(score)
    elif self.mask_nonlinear == "sigmoid":
        est_mask = torch.sigmoid(score)
    elif self.mask_nonlinear == "tanh":
        est_mask = torch.tanh(score)
    else:
        raise ValueError("Unsupported mask non-linear function")
    masks = est_mask.transpose(2, 3)  # [M, num_spk, K, N]
    masks = masks.unbind(dim=1)  # List[M, K, N]
    masked = [input * m for m in masks]
    others = OrderedDict(
        zip(["mask_spk{}".format(i + 1) for i in range(len(masks))], masks)
    )
    return masked, ilens, others