Additional code clean up