diff --git a/flaml/automl/__init__.py b/flaml/automl/__init__.py
index a517edf7e6..01a429d56d 100644
--- a/flaml/automl/__init__.py
+++ b/flaml/automl/__init__.py
@@ -1,5 +1,9 @@
-from flaml.automl.automl import AutoML, size
 from flaml.automl.logger import logger_formatter
-from flaml.automl.state import AutoMLState, SearchState
 
-__all__ = ["AutoML", "AutoMLState", "SearchState", "logger_formatter", "size"]
+try:
+    from flaml.automl.automl import AutoML, size
+    from flaml.automl.state import AutoMLState, SearchState
+
+    __all__ = ["AutoML", "AutoMLState", "SearchState", "logger_formatter", "size"]
+except ImportError:
+    __all__ = ["logger_formatter"]