diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
index 3f135e1c150..94e8ad43843 100644
--- a/.github/workflows/pytest.yml
+++ b/.github/workflows/pytest.yml
@@ -121,8 +121,7 @@ jobs:
 
           if [ "$PYTORCH" == "nightly" ]; then
             extra_index_url=https://download.pytorch.org/whl/nightly/cpu
-            # TODO: https://github.com/ludwig-ai/ludwig/issues/3395
-            pip install --pre torch==2.1.0.dev20230508+cpu torchtext torchvision torchaudio --index-url $extra_index_url
+            pip install --pre torch torchtext torchvision torchaudio --index-url $extra_index_url
           else
             extra_index_url=https://download.pytorch.org/whl/cpu
             pip install torch==$PYTORCH torchtext torchvision torchaudio --extra-index-url $extra_index_url