diff --git a/pyproject.toml b/pyproject.toml
index 13bcee2..6e59f8e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -48,7 +48,7 @@ dependencies = [ #TODO: can add optional dependencies
 'torch-geometric~=2.5.2',
 'torch-scatter~=2.1.2',
 'transformers~=4.42.4',
-'datasets~=3.2.0', #for HF datasets
+'datasets>=3.2,<3.6', #for HF datasets
 # 'causal-conv1d>=1.4.0', #For jamba/mamba
 # 'mamba-ssm~=2.2.4', #For jamba/mamba
 # 'flash-attn~=2.7.4.post1', #Flash Attention
diff --git a/requirements.txt b/requirements.txt
index be402e0..89e82cd 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -14,7 +14,7 @@ scipy~=1.15.2
 torch-geometric~=2.5.2
 torch-scatter~=2.1.2
 transformers~=4.42.4
-datasets~=3.2.0 #for HF datasets
+datasets>=3.2,<3.6 #for HF datasets
 # causal-conv1d>=1.4.0 #For jamba/mamba
 # mamba-ssm~=2.2.4 #For jamba/mamba
 # flash-attn~=2.7.4.post1 #Flash Attention