def zscore(input_x, dim=0):
    """Compute zscore over dim dimension.

    Parameters
    ----------
    input_x: Tensor
        Input tensor on which zscore will be applied
    dim: int, optional
        Dimension along which to operate. Default is 0.

    Returns
    -------
    zscore: Tensor
        A standardized tensor with the sale dimension as input_x
    """
    xt = input_x.transpose(0, dim)
    xt = (xt - xt.mean(0)) / xt.std(0)
    return xt.transpose(dim, 0)