import torch


def get_z_score(x: torch.Tensor, dim: int) -> torch.Tensor:
    """
    Standardize vector by removing the mean and scaling to unit variance

    Args:
        x: torch.Tensor
        dim: the dimension along which to compute the mean and std

    Return:
        The z-score, i.e. z = (x - mean) / std
    """
    std, mean = torch.std_mean(x, dim=dim, unbiased=True, keepdim=True)
    return (x-mean)/std