Fix bug in visualizing solution