rearrange

Author: casteryh

.github/workflows/unit_test.yaml

@@ -50,8 +50,11 @@ jobs:
           python -m pip install --upgrade pip
 
           # Install pytorch
-          python -m pip install torch==2.9.0.dev20250826 --extra-index-url https://download.pytorch.org/whl/nightly/cu126
 
+          python -m pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu126 --force-reinstall
+          # Install torchtitan
+
+          python -m pip install --pre torchtitan --index-url https://download.pytorch.org/whl/nightly/cu126 --force-reinstall
           # Install monarch
           python -m pip install monarch-no-torch==0.1.0.dev20250826 --find-links assets/ci
 
@@ -68,15 +71,6 @@ jobs:
           ssh-add torchstore_ssh_key
           python -m pip install git+ssh://[email protected]/meta-pytorch/torchstore.git
 
-          # torchtitan requires tomli>=1.1.0
-          pip install "tomli>=1.1.0"
-
-          # check if we have the correct version of tomli installed
-          pip freeze | grep tomli
-
-          # Install torchtitan
-          pip install --pre torchtitan --index-url https://download.pytorch.org/whl/nightly/cu126
-
           # Install dependencies
           python -m pip install --no-build-isolation -e ".[dev]"