Fixed a few issues with previous commit