Fix import problem when running in CLI.