Fix some bugs in the script