def predict_segmentation(logits):
    """
    Given the logits from a network, computing the segmentation by thresholding all values above 0 if `logits` has one
    channel, or computing the `argmax` along the channel axis otherwise, logits has shape `BCHW[D]`
    """
    if logits.shape[1] == 1:
        return (logits >= 0).int()  # for binary segmentation threshold on channel 0
    else:
        return logits.argmax(1).unsqueeze(1)