Added cleanup code