def normalize_latent_code(latent_code, adjust_norm=True):
    """Normalizes latent code.

    NOTE: The latent code will always be normalized along the last axis.
    Meanwhile, if `adjust_norm` is set as `True`, the norm of the result will be
    adjusted to `sqrt(latent_code.shape[-1])` in order to avoid too small value.

    Args:
        latent_code: The input latent code tensor to normalize.
        adjust_norm: Whether to adjust the norm of the output. (default: True)

    Returns:
        A normalized latent code tensor.

    Raises:
        TypeError: If the input `latent_code` is not with type `torch.Tensor`.
    """
    if not isinstance(latent_code, torch.Tensor):
        raise TypeError(f'Input latent code should be with type '
                        f'`torch.Tensor`, but `{type(latent_code)}` is '
                        f'received!')
    dim = latent_code.shape[-1]
    norm = latent_code.pow(2).sum(-1, keepdim=True).pow(0.5)
    out = latent_code / norm
    if adjust_norm:
        out = out * (dim ** 0.5)
    return out
Normalizes latent code.

    NOTE: The latent code will always be normalized along the last axis.
    Meanwhile, if `adjust_norm` is set as `True`, the norm of the result will be
    adjusted to `sqrt(latent_code.shape[-1])` in order to avoid too small value.

    Args:
        latent_code: The input latent code tensor to normalize.
        adjust_norm: Whether to adjust the norm of the output. (default: True)

    Returns:
        A normalized latent code tensor.

    Raises:
        TypeError: If the input `latent_code` is not with type `torch.Tensor`.
    
datasets/transforms.py