Fixed the code that didn't work from last commit