Fixed softmax confusion