Gain flattening