More code fixes

I'm stupid, I know.