Installation

First, install conda virtual environment for python >= 3.8.

conda install cmake ninja
# Run this command from the PyTorch directory after cloning the source code using the “Get the PyTorch Source“ section below
pip install -r requirements.txt

conda install intel::mkl-static intel::mkl-include
# CUDA only: Add LAPACK support for the GPU if needed
conda install -c pytorch magma-cuda110  # or the magma-cuda* that matches your CUDA version from https://anaconda.org/pytorch/repo

# Download source code.
git clone --recursive https://github.com/xwhzz/pytorch_stream_mask
cd pytorch

cd c10/libsmctrl
make libsmctrl.so
echo -e "\nexport LD_LIBRARY_PATH=$PWD:\$LD_LIBRARY_PATH" >> ~/.bashrc
source ~/.bashrc

cd ../..

export CMAKE_PREFIX_PATH=${CONDA_PREFIX:-"$(dirname $(which conda))/../"}
python setup.py develop

Test

python test_mask_stream/test_matmul.py