Fixed a stupid error with the grad calculation on neural net