More code fixes