Fixed bug in neural network