Register GPU kernel for TensorListElementShape

PiperOrigin-RevId: 258582149

Register GPU kernel for TensorListElementShape
PiperOrigin-RevId: 258582149
7eb4f192 · Eugene Zhulenev · TensorFlower Gardener · 5361b623 · 7eb4f192
隐藏空白更改
内联并排

Showing with 9 addition and 0 deletion

tensorflow/core/kernels/list_kernels.cc tensorflow/core/kernels/list_kernels.cc +9 -0

未找到文件。
--- a/tensorflow/core/kernels/list_kernels.cc
+++ b/tensorflow/core/kernels/list_kernels.cc
@@ -377,6 +377,15 @@ class TensorListElementShape : public OpKernel {
 REGISTER_KERNEL_BUILDER(Name("TensorListElementShape").Device(DEVICE_CPU),
                        TensorListElementShape);

+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+
+REGISTER_KERNEL_BUILDER(Name("TensorListElementShape")
+                            .Device(DEVICE_GPU)
+                            .HostMemory("element_shape"),
+                        TensorListElementShape);
+
+#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+
 class TensorListReserve : public OpKernel {
 public:
  explicit TensorListReserve(OpKernelConstruction* c) : OpKernel(c) {