Fix bug in cleanup