Small fixes for import errors