Minor update to the demo