Fix 413751357 updating version of depdendency which is lowering the JAX version (#167)

* Fix 413751357 updating orbax version which is lowering the JAX version

* Updating the right dependency

Author: parambole

requirements.txt

@@ -19,7 +19,7 @@ Pillow
 pylint
 pyink
 pytest==8.2.2
-tensorflow==2.17.0
+tensorflow>=2.17.0
 tensorflow-datasets>=4.9.6
 ruff>=0.1.5,<=0.2
 git+https://github.com/mlperf/logging.git