diff --git a/tests/e2e/mnist_pip_requirements.txt b/tests/e2e/mnist_pip_requirements.txt
index 3df95458..18548046 100644
--- a/tests/e2e/mnist_pip_requirements.txt
+++ b/tests/e2e/mnist_pip_requirements.txt
@@ -1,5 +1,5 @@
 --extra-index-url https://download.pytorch.org/whl/cu118
-torch==2.5.1
+torch==2.7.1+cu118
 torchvision==0.20.1
 pytorch_lightning==1.9.5
 torchmetrics==0.9.1