Update latest cuda image to include training-hub[cuda]

Author: Fiona-Waters

images/runtime/training/py312-cuda128-torch280/Dockerfile

@@ -82,10 +82,6 @@ RUN micropipenv install -- --no-cache-dir && \
     chmod -R g+w /opt/app-root/lib/python3.12/site-packages && \
     fix-permissions /opt/app-root -P
 
-# Install Flash Attention
-RUN pip install wheel
-RUN pip install --no-cache-dir flash-attn==2.8.3 --no-build-isolation
-
 # Upgrade NCCL to a more recent version and add Training Hub NVIDIA dependencies
 RUN pip install \
     nvidia-nccl-cu12==2.27.3 \

images/runtime/training/py312-cuda128-torch280/Pipfile

@@ -29,7 +29,7 @@ bitsandbytes = ">=0.45.3"
 liger-kernel = "==0.5.10"
 sentencepiece = "<0.3,>=0.1.99"
 tokenizers = "==0.21.4"
-training-hub = "==0.2.0"
+training-hub = {version = "==0.2.0", extras = ["cuda"]}
 trl = "==0.21.0"
 deepspeed = ">=0.14.3"
 async-timeout = "==4.0.3"