diff --git a/setup.py b/setup.py
index dbdf9745..28dd59bf 100644
--- a/setup.py
+++ b/setup.py
@@ -74,7 +74,7 @@ def read_version():
         "Programming Language :: Python :: 3.9",
     ],
     install_requires=[
-        "sagemaker-training>=4.3.0,<=4.8.3",
+        "sagemaker-training>=4.3.0,<=5.1.0",
         "numpy < 1.24",
         "scipy",
         "scikit-learn",