Changes to clean up code after merge