def METHOD_NAME(self, prob):
    # During training, torch.prod work on channel dimension.
    new_prob = torch.cat([torch.prod(1 - prob, dim=1, keepdim=True), prob],
                         1).clamp(1e-7, 1 - 1e-7)
    logits = torch.log((new_prob / (1 - new_prob)))
    return logits