Adding neural network fix