import torch


def normalize_rgb(raw_rgb, scaled_sigmoid):
    """Normalizes raw RGB values from the network.

    Args:
        raw_rgb: [R, S, 3] tf.float32. Raw RGB values.
        scaled_sigmoid: Whether to apply scaled sigmoid.

    Returns:
        rgb: [R, S, 3] tf.float32. Normalized RGB values.
    """
    # Extract RGB of each sample position along each ray.
    rgb = torch.sigmoid(raw_rgb)  # [R, S, 3]
    if scaled_sigmoid:
        rgb = 1.2 * (rgb - 0.5) + 0.5  # [R, S, 3]
    return rgb