Use int32 GPU kernel for StopGradient

PiperOrigin-RevId: 190204238

diff --git a/tensorflow/core/kernels/identity_op.cc b/tensorflow/core/kernels/identity_op.cc
index b5603fe..dffb4d7 100644 (file)
--- a/tensorflow/core/kernels/identity_op.cc
+++ b/tensorflow/core/kernels/identity_op.cc
@@ -129,6 +129,12 @@ REGISTER_GPU_KERNEL(Variant);
                               .HostMemory("output")           \
                               .TypeConstraint<type>("T"),     \
                           IdentityOp);                        \
+  REGISTER_KERNEL_BUILDER(Name("StopGradient")                \
+                              .Device(DEVICE_GPU)             \
+                              .HostMemory("input")            \
+                              .HostMemory("output")           \
+                              .TypeConstraint<type>("T"),     \
+                          IdentityOp);                        \
   REGISTER_KERNEL_BUILDER(Name("PlaceholderWithDefault")      \
                               .Device(DEVICE_GPU)             \
                               .HostMemory("input")            \