Fix gradient computation