import torch


def apply_gains_rgb(rgb, red_gains, blue_gains):
    """Applies white balance gains to a batch of RGB images."""
    red_gains = red_gains.squeeze()
    blue_gains = blue_gains.squeeze()
    red_gains = red_gains.unsqueeze(0) if len(red_gains.shape) == 0 else red_gains
    blue_gains = blue_gains.unsqueeze(0) if len(blue_gains.shape) == 0 else blue_gains
    rgb = rgb.permute(0, 2, 3, 1)  # Permute the image tensor to BxHxWxC format from BxCxHxW format
    green_gains = torch.ones_like(red_gains)
    gains = torch.stack([red_gains, green_gains, blue_gains], dim=-1)
    gains = gains[:, None, None, :]
    outs = rgb * gains
    outs = outs.permute(0, 3, 1, 2)  # Re-Permute the tensor back to BxCxHxW format
    return outs