diff --git a/colossalai/cli/__init__.py b/colossalai/cli/__init__.py
index e69de29bb..658e35e4c 100644
--- a/colossalai/cli/__init__.py
+++ b/colossalai/cli/__init__.py
@@ -0,0 +1,3 @@
+from .cli import cli
+
+__all__ = ['cli']
diff --git a/requirements/requirements.txt b/requirements/requirements.txt
index a05697b37..d60d4a81d 100644
--- a/requirements/requirements.txt
+++ b/requirements/requirements.txt
@@ -7,3 +7,4 @@ tensorboard
 packaging
 pre-commit
 rich
+click
diff --git a/setup.py b/setup.py
index 98d3d3e00..12b12c31d 100644
--- a/setup.py
+++ b/setup.py
@@ -215,7 +215,7 @@ setup(
     install_requires=fetch_requirements('requirements/requirements.txt'),
     entry_points='''
         [console_scripts]
-        colossal=colossalai.cli.cli:cli
+        colossal=colossalai.cli:cli
     ''',
     python_requires='>=3.7',
     classifiers=[