import torch


def apply_gains(bayer_images, red_gains, blue_gains):
  """Applies white balance gains to a batch of Bayer images."""
  red_gains = red_gains.squeeze(1)
  blue_gains= blue_gains.squeeze(1)
  bayer_images = bayer_images.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, green_gains, blue_gains], dim=-1)
  gains = gains[:, None, None, :]
  outs  = bayer_images * gains
  outs  = outs.permute(0, 3, 1, 2)  # Re-Permute the tensor back to BxCxHxW format
  return outs