Update gradients due to @pavlohuk comments