Use original TE from image

Author: Rohan-Bierneni

.github/workflows/UnitTests.yml

@@ -40,17 +40,12 @@ jobs:
 
       - name: Print dependencies
         run: |
-          pip uninstall -y transformer-engine transformer-engine-jax transformer-engine-cu12
-          pip install -U transformer-engine[jax]==2.6.0
-          # pip uninstall -y tensorflow
-          # pip install tensorflow-cpu
           pip freeze
 
       - name: Run MaxText Training
         run: |
           # This command is adapted from your DAG for a single-slice configuration.
           cd maxtext && \
-          pip install . --no-dependencies
 
           export XLA_PYTHON_CLIENT_MEM_FRACTION=0.65
           export TF_FORCE_GPU_ALLOW_GROWTH=true