Fix ending early bug