Fix gradient bug.

Does not work yet though.