Fix a potential bug

Just saw this when I read the code~

PS: Good job!