diff --git a/setup.py b/setup.py
index 2a2b5cc9cd..8c43a309c3 100644
--- a/setup.py
+++ b/setup.py
@@ -55,7 +55,7 @@
     'transformers>=4.36,<4.37',
     'mosaicml-streaming>=0.7.2,<0.8',
     'torch>=2.1,<2.1.1',
-    'datasets==2.15.0',
+    'datasets>=2.16,<2.17',
     'fsspec==2023.6.0',  # newer version results in a bug in datasets that duplicates data
     'sentencepiece==0.1.97',
     'einops==0.7.0',