diff options
Diffstat (limited to 'tensorflow/contrib/tensorrt/ops')
| -rw-r--r-- | tensorflow/contrib/tensorrt/ops/trt_engine_op.cc | 6 |
1 files changed, 6 insertions, 0 deletions
diff --git a/tensorflow/contrib/tensorrt/ops/trt_engine_op.cc b/tensorflow/contrib/tensorrt/ops/trt_engine_op.cc index 7139ff9618..fa72bce039 100644 --- a/tensorflow/contrib/tensorrt/ops/trt_engine_op.cc +++ b/tensorflow/contrib/tensorrt/ops/trt_engine_op.cc @@ -13,6 +13,9 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT + #include "tensorflow/core/framework/op.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/shape_inference.h" @@ -34,4 +37,7 @@ REGISTER_OP("TRTEngineOp") .Output("out_tensor: OutT") .SetShapeFn(shape_inference::TRTEngineOpShapeInference); +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA + } // namespace tensorflow |