Small update because of tree code changes