diff --git a/tensorflow_addons/custom_ops/text/cc/kernels/skip_gram_kernels.cc b/tensorflow_addons/custom_ops/text/cc/kernels/skip_gram_kernels.cc
index 4f1caa6da2..5b38d0665d 100644
--- a/tensorflow_addons/custom_ops/text/cc/kernels/skip_gram_kernels.cc
+++ b/tensorflow_addons/custom_ops/text/cc/kernels/skip_gram_kernels.cc
@@ -133,7 +133,7 @@ class SkipGramGenerateCandidatesOp : public OpKernel {
                               .TypeConstraint<type>("T"),           \
                           SkipGramGenerateCandidatesOp<type>)
 
-REGISTER_KERNEL(string);
+REGISTER_KERNEL(tstring);
 REGISTER_KERNEL(int64);
 REGISTER_KERNEL(int32);
 REGISTER_KERNEL(int16);
@@ -141,4 +141,4 @@ REGISTER_KERNEL(int16);
 #undef REGISTER_KERNEL
 
 }  // end namespace addons
-}  // namespace tensorflow
\ No newline at end of file
+}  // namespace tensorflow