install_fp.sh

APEX_DIR=apex
N_PROC=8
CUDA_VERSION=11.4.0

# Install packages
mamba install -c conda-forge \
      cudatoolkit=11.4.0 cudatoolkit-dev=11.4.0 pytorch-gpu=1.9.0 \
      scipy numpy matplotlib tensorboardx tensorboard inflect \
      gcc=9.4.0 gxx=9.4.0

export CUDA_HOME=$CONDA_PREFIX

# Install dll logger
pip install git+https://github.com/NVIDIA/dllogger@v0.1.0#egg=dllogger

# Install apex
git clone https://github.com/NVIDIA/apex.git $APEX_DIR
(
    cd $APEX_DIR
    export LIBRARY_PATH=$CONDA_PREFIX/lib
    MAX_JOBS=$N_PROC pip install -v --disable-pip-version-check --no-cache-dir \
                     --global-option="--cpp_ext" --global-option="--cuda_ext" \
                     ./
)