diff options
Diffstat (limited to 'tensorflow/core/kernels/stage_op.cc')
| -rw-r--r-- | tensorflow/core/kernels/stage_op.cc | 6 |
1 files changed, 6 insertions, 0 deletions
diff --git a/tensorflow/core/kernels/stage_op.cc b/tensorflow/core/kernels/stage_op.cc index c18b992ea1..161ba89212 100644 --- a/tensorflow/core/kernels/stage_op.cc +++ b/tensorflow/core/kernels/stage_op.cc @@ -99,6 +99,9 @@ REGISTER_KERNEL_BUILDER(Name("Stage").Device(DEVICE_CPU), StageOp); #if GOOGLE_CUDA REGISTER_KERNEL_BUILDER(Name("Stage").Device(DEVICE_GPU), StageOp); #endif +#ifdef TENSORFLOW_USE_SYCL +REGISTER_KERNEL_BUILDER(Name("Stage").Device(DEVICE_SYCL), StageOp); +#endif // TENSORFLOW_USE_SYCL class UnstageOp : public OpKernel { public: @@ -126,5 +129,8 @@ REGISTER_KERNEL_BUILDER(Name("Unstage").Device(DEVICE_CPU), UnstageOp); #if GOOGLE_CUDA REGISTER_KERNEL_BUILDER(Name("Unstage").Device(DEVICE_GPU), UnstageOp); #endif +#ifdef TENSORFLOW_USE_SYCL +REGISTER_KERNEL_BUILDER(Name("Unstage").Device(DEVICE_SYCL), UnstageOp); +#endif // TENSORFLOW_USE_SYCL } // namespace tensorflow |