diff options
Diffstat (limited to 'tensorflow/contrib/memory_stats/kernels/memory_stats_ops.cc')
| -rw-r--r-- | tensorflow/contrib/memory_stats/kernels/memory_stats_ops.cc | 11 |
1 files changed, 11 insertions, 0 deletions
diff --git a/tensorflow/contrib/memory_stats/kernels/memory_stats_ops.cc b/tensorflow/contrib/memory_stats/kernels/memory_stats_ops.cc index 23a682fb90..3b88535dce 100644 --- a/tensorflow/contrib/memory_stats/kernels/memory_stats_ops.cc +++ b/tensorflow/contrib/memory_stats/kernels/memory_stats_ops.cc @@ -57,6 +57,11 @@ REGISTER_KERNEL_BUILDER(Name("BytesLimit").Device(DEVICE_CPU), BytesLimitOp); REGISTER_KERNEL_BUILDER(Name("BytesLimit").Device(DEVICE_GPU).HostMemory("out"), BytesLimitOp); +#ifdef TENSORFLOW_USE_SYCL +REGISTER_KERNEL_BUILDER(Name("BytesLimit").Device(DEVICE_SYCL).HostMemory("out"), + BytesLimitOp); +#endif // TENSORFLOW_USE_SYCL + // Op that measures the peak memory in bytes. class MaxBytesInUseOp : public MemoryStatsOp { public: @@ -76,4 +81,10 @@ REGISTER_KERNEL_BUILDER( Name("MaxBytesInUse").Device(DEVICE_GPU).HostMemory("out"), MaxBytesInUseOp); +#ifdef TENSORFLOW_USE_SYCL +REGISTER_KERNEL_BUILDER( + Name("MaxBytesInUse").Device(DEVICE_SYCL).HostMemory("out"), + MaxBytesInUseOp); +#endif // TENSORFLOW_USE_SYCL + } // namespace tensorflow |