diff --git a/.github/workflows/unit_test_4gpu.yaml b/.github/workflows/unit_test_4gpu.yaml
index 669975f3..32852a35 100644
--- a/.github/workflows/unit_test_4gpu.yaml
+++ b/.github/workflows/unit_test_4gpu.yaml
@@ -34,6 +34,7 @@ jobs:
       - name: Install dependencies
         run: |
           pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121
+          pip3 install --pre torchdata --index-url https://download.pytorch.org/whl/nightly/cpu
           python -m pip install -r requirements.txt
           python -m pip install -r dev-requirements.txt
           python -m pip install -e .
diff --git a/.github/workflows/unit_test_cpu.yaml b/.github/workflows/unit_test_cpu.yaml
index 354182ac..097a16a6 100644
--- a/.github/workflows/unit_test_cpu.yaml
+++ b/.github/workflows/unit_test_cpu.yaml
@@ -34,6 +34,7 @@ jobs:
       - name: Install dependencies
         run: |
           pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121
+          pip3 install --pre torchdata --index-url https://download.pytorch.org/whl/nightly/cpu
           python -m pip install -r requirements.txt
           python -m pip install -r dev-requirements.txt
           python -m pip install -e .
diff --git a/requirements.txt b/requirements.txt
index 24f46e8f..b82120a6 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -5,5 +5,3 @@ tensorboard
 sentencepiece
 tiktoken
 blobfile
--i https://download.pytorch.org/whl/nightly/cpu
-torchdata