Instructions to build pytorch with CRAY MPICH

build_pytorch.sh

# Make a conda env
conda create -y --prefix $INSTALL_DIR python=3.6
conda activate $INSTALL_DIR
conda install -y numpy pyyaml mkl mkl-include setuptools cmake cffi typing
conda install -y -c mingfeima mkldnn

# Configure the build using Cray compiler wrappers cc, CC
export CMAKE_PREFIX_PATH="$(dirname $(which conda))/../"
export CXX=CC
export CC=cc
export NO_CUDA=1
export CRAYPE_LINK_TYPE=dynamic
export WITH_DISTRIBUTED=1
export CFLAGS="-L/opt/cray/pe/mpt/7.7.0/gni/mpich-gnu/51/lib -lmpich_gnu_51_mt"
export CXXFLAGS="-L/opt/cray/pe/mpt/7.7.0/gni/mpich-gnu/51/lib -lmpichcxx_gnu_51_mt"

# Download pytorch
git clone --recursive --branch $PYTORCH_VERSION https://github.com/pytorch/pytorch

# Build PyTorch
cd pytorch
python setup.py install