File tree Expand file tree Collapse file tree 3 files changed +8
-7
lines changed Expand file tree Collapse file tree 3 files changed +8
-7
lines changed Original file line number Diff line number Diff line change @@ -48,9 +48,7 @@ HOST_MODEL_DIR=${HOST_MODEL_DIR:="${HOST_BUILD_DIR}/${TRITON_VERSION}"}
48
48
HOST_SOURCE_DIR=$HOST_BUILD_DIR /gen_srcdir
49
49
50
50
# Set CI specific parameters
51
- DOCKER_GPU_ARGS=" ${DOCKER_GPU_ARGS:= " --gpus device=$CUDA_DEVICE " } "
52
- [[ $RUNNER_GPUS =~ ^[0-9] ]] && DOCKER_GPU_ARGS=$( eval $NV_DOCKER_ARGS )
53
-
51
+ DOCKER_GPU_ARGS=${DOCKER_GPU_ARGS:- $([[ $RUNNER_GPUS =~ ^[0-9] ]] && eval $NV_DOCKER_ARGS || echo " --gpus device=$CUDA_DEVICE " )}
54
52
55
53
# Set model output directories
56
54
Original file line number Diff line number Diff line change @@ -44,7 +44,7 @@ PYTORCH_IMAGE=${PYTORCH_IMAGE:=nvcr.io/nvidia/pytorch:$NVIDIA_UPSTREAM_VERSION-p
44
44
45
45
CUDA_DEVICE=${NV_GPU:= 0}
46
46
47
- [[ $RUNNER_GPUS =~ ^[0-9] ]] && DOCKER_GPU_ARGS= $( eval $NV_DOCKER_ARGS ) || DOCKER_GPU_ARGS= " --gpus device=$CUDA_DEVICE "
47
+ DOCKER_GPU_ARGS= ${DOCKER_GPU_ARGS :- $( [[ $RUNNER_GPUS =~ ^[0-9] ]] && eval $NV_DOCKER_ARGS || echo " --gpus device=$CUDA_DEVICE " )}
48
48
49
49
# ##
50
50
HOST_BUILD_DIR=${HOST_BUILD_DIR:=/ tmp}
Original file line number Diff line number Diff line change @@ -63,7 +63,8 @@ TENSORFLOW_IMAGE=${TENSORFLOW_IMAGE:=nvcr.io/nvidia/tensorflow:$TRITON_VERSION-t
63
63
TENSORRT_IMAGE=${TENSORRT_IMAGE:= nvcr.io/ nvidia/ tensorrt: $TRITON_VERSION -py3}
64
64
CUDA_DEVICE=${NV_GPU:= 0}
65
65
66
- [[ $RUNNER_GPUS =~ ^[0-9] ]] && DOCKER_GPU_ARGS=$( eval $NV_DOCKER_ARGS ) || DOCKER_GPU_ARGS=" --gpus device=$CUDA_DEVICE "
66
+ DOCKER_GPU_ARGS=${DOCKER_GPU_ARGS:- $([[ $RUNNER_GPUS =~ ^[0-9] ]] && eval $NV_DOCKER_ARGS || echo " --gpus device=$CUDA_DEVICE " )}
67
+ MODEL_TYPE=${MODEL_TYPE:- " " }
67
68
68
69
# ##
69
70
HOST_BUILD_DIR=${HOST_BUILD_DIR:=/ tmp}
@@ -360,8 +361,10 @@ python3 $SRCDIR/gen_qa_implicit_models.py --libtorch --variable --models_dir=$VA
360
361
chmod -R 777 $VARIMPLICITSEQDESTDIR
361
362
python3 $SRCDIR /gen_qa_dyna_sequence_models.py --libtorch --models_dir=$DYNASEQDESTDIR
362
363
chmod -R 777 $DYNASEQDESTDIR
363
- python3 $SRCDIR /gen_qa_torchtrt_models.py --models_dir=$TORCHTRTDESTDIR
364
- chmod -R 777 $TORCHTRTDESTDIR
364
+ if [ -z "$MODEL_TYPE " ] || [ "$MODEL_TYPE " != "igpu" ]; then
365
+ python3 $SRCDIR /gen_qa_torchtrt_models.py --models_dir=$TORCHTRTDESTDIR
366
+ chmod -R 777 $TORCHTRTDESTDIR
367
+ fi
365
368
python3 $SRCDIR /gen_qa_ragged_models.py --libtorch --models_dir=$RAGGEDDESTDIR
366
369
chmod -R 777 $RAGGEDDESTDIR
367
370
# Export torchvision image models to ONNX
You can’t perform that action at this time.
0 commit comments