fixing bug in stochastic gradient