[1.x] Enable CUDA Graphs for TRT (#19184)

Signed-off-by: Serge Panev <[email protected]>
apache · Sep 19, 2020 · a35d568 · a35d568
1 parent 0496690
commit a35d568
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/operator/subgraph/tensorrt/tensorrt.cc b/src/operator/subgraph/tensorrt/tensorrt.cc
@@ -370,6 +370,11 @@ NNVM_REGISTER_OP(_TensorRT)
     .set_attr<nnvm::FListInputNames>("FListInputNames", TRTListInputNames)
     .set_attr<nnvm::FListOutputNames>("FListOutputNames", DefaultSubgraphOpListOutputs)
     .set_attr<FCreateOpState>("FCreateOpState", TRTCreateState)
+    .set_attr<FIsCUDAGraphsCompatible>("FIsCUDAGraphsCompatible",
+        [](const NodeAttrs& attrs, const bool) {
+          const TRTParam& param = nnvm::get<TRTParam>(attrs.parsed);
+          return !param.int8_mode;
+        })
     .set_attr<FInferStorageType>("FInferStorageType", TRTInferStorageType);
 
 MXNET_REGISTER_SUBGRAPH_BACKEND(TensorRT);