Fixed a few issues in last commit