Small fixes from last commit