Minor gradient update fix