Specify path for PyTorch model extension library (#7025)

triton-inference-server · Mar 26, 2024 · 1441f45 · 1441f45
1 parent 5516ee1
commit 1441f45
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/qa/common/gen_qa_custom_ops b/qa/common/gen_qa_custom_ops
@@ -124,12 +124,15 @@ if [ $? -ne 0 ]; then
 fi
 
 # PyTorch
+
 cat >$HOST_SRCDIR/$PYTSCRIPT <<EOF
 #!/bin/bash -x
 nvidia-smi -L || true
 nvidia-smi || true
 set -e
+export TORCH_EXTENSIONS_DIR="/root/.cache/torch_extensions/"
 python3 $SRCDIR/gen_qa_custom_ops_models.py --libtorch --models_dir=$DESTDIR
+cp \${TORCH_EXTENSIONS_DIR}/custom_modulo/custom_modulo.so $DESTDIR/libtorch_modulo/.
 chmod -R 777 $DESTDIR
 EOF