Cleanup bug fix