setup.sh

#!/bin/bash
set -e

# Install flash-attn with no-build-isolation
pip install flash-attn==2.7.3 --no-build-isolation

# Uninstall apex if it exists (as in your original script)
pip uninstall apex -y || true