don't build images w cuda 130 since we don't have flash attention wheels (#3341)

.github/workflows/main.yml

@@ -31,11 +31,11 @@ jobs:
             python_version: "3.11"
             pytorch: 2.9.1
             axolotl_extras:
-          - cuda: 130
+#          - cuda: 130
-            cuda_version: 13.0.0
+#            cuda_version: 13.0.0
-            python_version: "3.11"
+#            python_version: "3.11"
-            pytorch: 2.9.1
+#            pytorch: 2.9.1
-            axolotl_extras:
+#            axolotl_extras:
     runs-on: axolotl-gpu-runner
     steps:
       - name: Checkout
@@ -98,11 +98,11 @@ jobs:
             python_version: "3.11"
             pytorch: 2.9.1
             axolotl_extras:
-          - cuda: 130
+#          - cuda: 130
-            cuda_version: 13.0.0
+#            cuda_version: 13.0.0
-            python_version: "3.11"
+#            python_version: "3.11"
-            pytorch: 2.9.1
+#            pytorch: 2.9.1
-            axolotl_extras:
+#            axolotl_extras:
     runs-on: axolotl-gpu-runner
     steps:
       - name: Checkout

docker/Dockerfile-base

@@ -51,7 +51,7 @@ RUN git lfs install --skip-repo && \
     pip3 install -U --no-cache-dir pydantic==1.10.10 && \
     pip3 cache purge
 
-RUN if [ "$PYTORCH_VERSION" = "2.9.1" ] && [ "$CUDA" = "128" ] ; then \
+RUN if [ "$PYTORCH_VERSION" =~ ^2\.9\.[0-9]+$ ] && [ "$CUDA" = "128" ] ; then \
         wget https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.4.17/flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \
         pip3 install --no-cache-dir flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \
         rm flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \