Use 32bit indexing to pad tensors. This more than doubles the performance of

the corresponding TensorFlow operation.
Change: 116831272

1 files changed, 5 insertions, 1 deletions

diff --git a/tensorflow/core/kernels/pad_op.h b/tensorflow/core/kernels/pad_op.h
index cdcd248bd9..fcc808b01f 100644
--- a/tensorflow/core/kernels/pad_op.h
+++ b/tensorflow/core/kernels/pad_op.h
@@ -32,7 +32,11 @@ struct Pad {
   void operator()(const Device& d, typename TTypes<T, Dims>::Tensor output,
                   typename TTypes<T, Dims>::ConstTensor input,
                   Eigen::array<std::pair<int32, int32>, Dims> paddings) {
-    output.device(d) = input.pad(paddings);
+    if (Eigen::internal::is_same<Device, Eigen::GpuDevice>::value) {
+      To32Bit(output).device(d) = To32Bit(input).pad(paddings);
+    } else {
+      output.device(d) = input.pad(paddings);
+    }
   }
 };