diff options
| author |  A. Unique TensorFlower <gardener@tensorflow.org> | 2018-08-15 16:19:19 -0700 |
|---|---|---|
| committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-08-15 16:26:59 -0700 |
| commit | 698c9a1b6344af44e92eb7fbd2a656d7a19b43d3 (patch) | |
| tree | 9a619109b984551aabedd3eb2861c8c0775c9a00 /tensorflow/contrib/tensorrt | |
| parent | 82820a6b6263525781a16c9d5fbf745d3f24ebd7 (diff) | |
Make sure tensorrt and nccl kernels compile with newer versions of Eigen.
PiperOrigin-RevId: 208902416
Diffstat (limited to 'tensorflow/contrib/tensorrt')
| -rw-r--r-- | tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_kernel.cu.cc | 9 |
1 files changed, 5 insertions, 4 deletions
diff --git a/tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_kernel.cu.cc b/tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_kernel.cu.cc index 2de7973750..11335d7da6 100644 --- a/tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_kernel.cu.cc +++ b/tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_kernel.cu.cc @@ -13,14 +13,15 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT + #include "tensorflow/contrib/tensorrt/custom_plugin_examples/inc_op_kernel.h" #include <vector> +#define EIGEN_USE_GPU #include "tensorflow/core/framework/op_kernel.h" - -#if GOOGLE_CUDA -#if GOOGLE_TENSORRT #include "cuda/include/cuda_runtime_api.h" #include "tensorflow/core/platform/stream_executor.h" @@ -80,5 +81,5 @@ REGISTER_KERNEL_BUILDER(Name("IncPluginTRT").Device(DEVICE_GPU), IncPluginTRT); } // namespace tensorrt } // namespace tensorflow -#endif // GOOGLE_CUDA #endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA |