From fe020fccbb7a08f386d351877937dc9cdb4554d5 Mon Sep 17 00:00:00 2001 From: Trevor Morris Date: Mon, 11 Mar 2019 14:50:02 -0700 Subject: [PATCH] Store an empty engine in the case that dynamic engine creation fails. The next time the same input shape appears, we will retrieve the empty engine which tells us to not rebuild the same failing engine again. --- tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc index 0f800d7cf26..e612d311bce 100644 --- a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc +++ b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc @@ -633,6 +633,9 @@ EngineContext* TRTEngineOp::GetEngine( } LOG(WARNING) << "Engine creation for batch size " << batch_size << " failed " << status; + // Store an empty engine here so we don't try to build the same engine + // again. + cache.emplace(engine_input_shapes, empty_context); return &empty_context; } VLOG(1) << "Conversion is done";