Solved a slight bug in my last commit