Updates to demo code