Fixed call to build_model_bax_schwarz

Previously I had forgotten to change it from build_model_bax_heat