def permute_data(x, pe_dir):
    """
    This function permutes data of the tensor x of size (batch, n_channels, 1st_size, ..., nth_size)
    so that the PE direction become the first dimension (except the batch and n_channels dimensions).
    :param x: the tensor of size (batch, n_channels, 1st_size, ..., nth_size)
    :param pe_dir: the number indicating the PE direction of the image of size (1st_size, ..., nth_size)
    :return: a tensor of size (batch, n_channels, pe_direction_th_size, ..., 1st_size, ..., nth_size)
    """
    # set the permutation direction so that the pe_dir and the first directions are swapped
    per_seq = list(range(len(x.shape) - 2))
    per_seq[0] = pe_dir
    per_seq[pe_dir] = 0
    per_seq = [k + 2 for k in per_seq]

    # permute the data
    if len(x.shape) - 2 == 2:  # for 2D images
        x = x.permute(0, 1, per_seq[0], per_seq[1])
    elif len(x.shape) - 2 == 3:  # for 3D images
        x = x.permute(0, 1, per_seq[0], per_seq[1], per_seq[2])
    return x