softmax bug fix