Following is the script that I used to compile pytorch from source.
mkdir $HOME/.openmpi/
wget https://www.open-mpi.org/software/ompi/v3.0/downloads/openmpi-3.0.0.tar.gz
gunzip -c openmpi-3.0.0.tar.gz | tar xf - \
    && cd openmpi-3.0.0 \
    && ./configure --prefix=$HOME/.openmpi/ --with-cuda=/usr/local/cuda-10.1 \
    && make all install
export PATH=$HOME/.openmpi/bin:$PATH
export LD_LIBRARY_PATH=$HOME/.openmpi/lib:$LD_LIBRARY_PATH

######### add cuda directories 

export CUDA_NVCC_EXECUTABLE="/usr/local/cuda-10.1/bin/nvcc"
export CUDA_HOME="/usr/local/cuda-10.1/"
export CUDNN_INCLUDE_PATH="/usr/local/cuda-10.1/include/"
export CUDNN_LIBRARY_PATH="/usr/local/cuda-10.1/lib64/"
export LIBRARY_PATH="/usr/local/cuda-10.1/lib64"
export USE_CUDA=1 USE_CUDNN=1 USE_MKLDNN=1 MAX_JOBS=80


export ENV_NAME=pt16
conda create --name $ENV_NAME python=3.7.3 numpy=1.16.3 
conda activate $ENV_NAME 
conda install numpy pyyaml mkl mkl-include setuptools cmake cffi typing_extensions

# use cuda101 to match with the system 
conda install -c pytorch magma-cuda101 
export CMAKE_PREFIX_PATH="$HOME/anaconda3/envs/$ENV_NAME"
cd ~/anaconda3/envs/$ENV_NAME/compiler_compat
mv ld ld-old
sudo apt-get install libomp-dev
cd ~
git clone --recursive https://github.com/pytorch/pytorch pytorch && \
cd pytorch && \
git checkout v1.6.0 --recurse-submodules && \
python setup.py install | tee 2>&1 > compile.log