Bug fixing for demo