diff --git a/setup.py b/setup.py
index 196591cfe0..b063b08c69 100644
--- a/setup.py
+++ b/setup.py
@@ -29,7 +29,7 @@
 INSTALL_REQUIRE = [
     "torch>=1.11",
     "transformers>=4.36.0,<4.43.0",
-    "optimum @ git+https://github.com/huggingface/optimum.git@transformer-442",
+    "optimum @ git+https://github.com/huggingface/optimum.git@fix-sentence-trfs",
     "datasets>=1.4.0",
     "sentencepiece",
     "setuptools",