Fix softmax input_dim
Fix softmax input_dim

Otherwise when shared softmax is used the input_dim is 0 by default.

PiperOrigin-RevId: 304238122

tensorflow/lingvo,tensorflow/lingvo,tensorflow/lingvo,tensorflow/lingvo