Updating to work with new solver