Fix bugs after demo