diff --git a/src/operator/subgraph/tensorrt/tensorrt.cc b/src/operator/subgraph/tensorrt/tensorrt.cc index 1c60d026afd0..71d096edfc81 100644 --- a/src/operator/subgraph/tensorrt/tensorrt.cc +++ b/src/operator/subgraph/tensorrt/tensorrt.cc @@ -370,6 +370,11 @@ NNVM_REGISTER_OP(_TensorRT) .set_attr("FListInputNames", TRTListInputNames) .set_attr("FListOutputNames", DefaultSubgraphOpListOutputs) .set_attr("FCreateOpState", TRTCreateState) + .set_attr("FIsCUDAGraphsCompatible", + [](const NodeAttrs& attrs, const bool) { + const TRTParam& param = nnvm::get(attrs.parsed); + return !param.int8_mode; + }) .set_attr("FInferStorageType", TRTInferStorageType); MXNET_REGISTER_SUBGRAPH_BACKEND(TensorRT);