Fix end game bug