git clone https://github.com/ACEsuit/mace.git
pip install ./mace

git clone --branch=mace --depth=1 https://github.com/ACEsuit/lammps
wget https://download.pytorch.org/libtorch/cu121/libtorch-shared-with-deps-2.2.0%2Bcu121.zip
unzip libtorch-shared-with-deps-2.2.0+cu121.zip
rm libtorch-shared-with-deps-2.2.0+cu121.zip
mv libtorch libtorch-gpu

conda config --add channels defaults
conda config --add channels conda-forge
conda install mkl mkl-include -y

# Set CUDA environment variables for CUDA 12.1

export MKL_INCLUDE_DIR=$CONDA_PREFIX/include
export MKL_LIBRARY_DIR=$CONDA_PREFIX/lib

module load cuda-12.1 gcc

cmake \
    -D CMAKE_BUILD_TYPE=Release \
    -D CMAKE_INSTALL_PREFIX=$(pwd) \
    -D CMAKE_CXX_STANDARD=17 \
    -D CMAKE_CXX_STANDARD_REQUIRED=ON \
    -DMKL_INCLUDE_DIR=$MKL_INCLUDE_DIR \
    -DMKL_LIBRARY_DIR=$MKL_LIBRARY_DIR \
    -D BUILD_MPI=yes \
    -D BUILD_SHARED_LIBS=yes \
    -D PKG_KOKKOS=yes \
    -D Kokkos_ARCH_SKX=yes \
    -D Kokkos_ARCH_AMPERE80=ON \
    -D Kokkos_ENABLE_CUDA=yes \
    -D CMAKE_CXX_COMPILER=$(pwd)/../lib/kokkos/bin/nvcc_wrapper \
    -D PKG_ML-MACE=yes \
    -D CMAKE_PREFIX_PATH=$(pwd)/../../libtorch-gpu \
    ../cmake

make -j 32