fix the bug of environment variables

Dockerfile

@@ -34,9 +34,8 @@ ENV HOME=/home/user \
 
 RUN pip3 install torch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 --index-url https://download.pytorch.org/whl/cu121
 RUN pip3 install --no-cache-dir --upgrade -r /code/requirements.txt
-RUN python /code/get_flash_attn.py > /code/url.txt
-RUN source /code/url.txt
-RUN pip3 install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.2.post1/${FLASH_ATTN_VER} --no-cache-dir
+RUN FLASH_ATTN_VER=$(python /code/get_flash_attn.py) && \
+    pip3 install "https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.2.post1/${FLASH_ATTN_VER}" --no-cache-dir
 
 # Set the working directory to the user's home directory
 WORKDIR $HOME/app

get_flash_attn.py

@@ -54,6 +54,6 @@ def generate_flash_attn_filename(flash_attn_version="2.7.2.post1"):
 if __name__ == "__main__":
     try:
         filename = generate_flash_attn_filename()
-        print(f"FLASH_ATTN_VER={filename}")
+        print(f"{filename}")
     except Exception as e:
         print("Error generating filename:", e)