SoftMaxTree fixes

This PR includes some fixes to SoftMaxTree.