Update gradients due to @auvalent comments