Fix some bugs in the last commit