Fixed my own buggy code