Additional code fixes