Update UnitTests.yml

Author: parambole

.github/workflows/UnitTests.yml

@@ -31,7 +31,7 @@ jobs:
     # IMPORTANT: Replace with the label for your specific GPU runner if different
     runs-on: ["linux-x86-a3-megagpu-h100-8gpu"]
     container:
-      image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/gpu:jax0.7.2-cuda12.9-rev1:latest
+      image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/gpu:jax0.7.2-cuda12.9-rev1
 
     steps:
       - name: Checkout Repository
@@ -58,4 +58,4 @@ jobs:
             output_dir="/tmp/flux-output/" \
             max_train_steps=5 \
             hardware=gpu \
-            attention="cudnn_flash_te"
+            attention="cudnn_flash_te"