try without te

Author: Rohan-Bierneni

.github/workflows/UnitTests.yml

@@ -69,7 +69,7 @@ jobs:
     # IMPORTANT: Replace with the label for your runner (e.g., v5p-8)
     runs-on: ["linux-x86-a3-megagpu-h100-8gpu"]
     container:
-      image: gcr.io/tpu-prod-env-multipod/maxdiffusion_stable_stack_candidate:jax0.7.2_cuda13_te2.6.0
+      image: gcr.io/tpu-prod-env-multipod/maxdiffusion_stable_stack_candidate:
     steps:
       - name: Checkout Repository
         uses: actions/checkout@v4
@@ -89,7 +89,7 @@ jobs:
 
       - name: Print dependencies
         run: |
-          # pip uninstall -y transformer-engine transformer-engine-jax transformer-engine-cu12
+          pip uninstall -y transformer-engine transformer-engine-jax transformer-engine-cu12
           # pip install transformer_engine[jax]==2.4.0
           # pip install -U transformer-engine[jax]==2.6.0
           # pip uninstall -y transformer-engine-cu12
@@ -108,7 +108,7 @@ jobs:
             train_text_encoder=false \
             cache_latents_text_encoder_outputs=true \
             per_device_batch_size=1 \
-            attention=dot_product \
+            attention=cudnn_flash_te \
             activations_dtype=bfloat16 \
             weights_dtype=bfloat16 \
             max_train_steps=200 \