diff --git a/catx/VERSION b/catx/VERSION
index b1e80bb..845639e 100644
--- a/catx/VERSION
+++ b/catx/VERSION
@@ -1 +1 @@
-0.1.3
+0.1.4
diff --git a/catx/type_defs.py b/catx/type_defs.py
index 074519e..e4e0bc6 100644
--- a/catx/type_defs.py
+++ b/catx/type_defs.py
@@ -1,6 +1,9 @@
-from typing import TypeAlias
+try:
+    from typing import TypeAlias
+except ImportError:
+    from typing_extensions import TypeAlias
 
-from chex import ArrayNumpy, Array
+from chex import Array, ArrayNumpy
 
 Actions: TypeAlias = ArrayNumpy
 JaxActions: TypeAlias = Array
diff --git a/requirements.txt b/requirements.txt
index 6aa3f62..5102ba1 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
+chex
 dm-haiku
 matplotlib
 numpy
diff --git a/setup.py b/setup.py
index 137051c..4194424 100644
--- a/setup.py
+++ b/setup.py
@@ -39,6 +39,7 @@ def read_requirements(path: str) -> List[str]:
     long_description_content_type="text/markdown",
     author="InstaDeep",
     packages=find_packages(),
+    python_requires=">=3.7",
     install_requires=read_requirements("requirements.txt"),
     entry_points={"console_scripts": ["project_name = project_name.__main__:main"]},
     extras_require={