Fixed some bugs with the latest changes.