fixing demo code