Check jax.devices output

Author: Rohan-Bierneni

.github/workflows/UnitTests.yml

@@ -98,23 +98,27 @@ jobs:
           # pip install tensorflow-cpu
           pip freeze
 
-      - name: Run MaxDiffusion Training
+      - name: Check per_device_batch_size
         run: |
-          # This command is adapted from your DAG for a single-slice configuration.
-          NVTE_FUSED_ATTN=1 pip install . && \
-            python -m src.maxdiffusion.train_sdxl src/maxdiffusion/configs/base_xl.yml \
-            hardware=gpu \
-            train_new_unet=true \
-            train_text_encoder=false \
-            cache_latents_text_encoder_outputs=true \
-            per_device_batch_size=1 \
-            attention=dot_product \
-            activations_dtype=bfloat16 \
-            weights_dtype=bfloat16 \
-            max_train_steps=200 \
-            enable_profiler=True \
-            run_name=1slice-VGpuVersion.XPK_H100_a3-maxdiffusion-jax-stable-stack-2025-09-26-04-12-02 \
-            output_dir=gs://rbierneni-multipod-dev/${{ github.run_id }}
+          python -c "import jax; print(jax.devices())"
+
+      # - name: Run MaxDiffusion Training
+      #   run: |
+      #     # This command is adapted from your DAG for a single-slice configuration.
+      #     NVTE_FUSED_ATTN=1 pip install . && \
+      #       python -m src.maxdiffusion.train_sdxl src/maxdiffusion/configs/base_xl.yml \
+      #       hardware=gpu \
+      #       train_new_unet=true \
+      #       train_text_encoder=false \
+      #       cache_latents_text_encoder_outputs=true \
+      #       per_device_batch_size=1 \
+      #       attention=dot_product \
+      #       activations_dtype=bfloat16 \
+      #       weights_dtype=bfloat16 \
+      #       max_train_steps=200 \
+      #       enable_profiler=True \
+      #       run_name=1slice-VGpuVersion.XPK_H100_a3-maxdiffusion-jax-stable-stack-2025-09-26-04-12-02 \
+      #       output_dir=gs://rbierneni-multipod-dev/${{ github.run_id }}
 
 # jobs:
 #   build: