Fix for gradient bug