[Embedding] Enable EmbeddingVariable python UT in cibuild. (#643)

Author: candyzone

cibuild/gpu-ut/gpu-python-ut.sh

@@ -107,8 +107,6 @@ export TF_BUILD_BAZEL_TARGET="$TF_ALL_TARGETS "\
 "-//tensorflow/python/keras:convolutional_test "\
 "-//tensorflow/python/keras:lstm_v2_test "\
 "-//tensorflow/python/keras:lstm_v2_test_gpu "\
-"-//tensorflow/python:embedding_variable_ops_gpu_test "\
-"-//tensorflow/python:embedding_variable_ops_gpu_test_gpu "\
 "-//tensorflow/python/kernel_tests:normalize_op_test "\
 "-//tensorflow/python/kernel_tests:svd_op_test "
 

tensorflow/core/kernels/kv_variable_ops.cc

@@ -95,7 +95,8 @@ class KvVariableShapeOp : public OpKernel {
       Name("KvVariableShape").Device(DEVICE_##dev)                    \
                              .TypeConstraint<type>("out_type")        \
                              .TypeConstraint<ktype>("Tkeys")          \
-                             .TypeConstraint<vtype>("dtype"),         \
+                             .TypeConstraint<vtype>("dtype")          \
+                             .HostMemory("output"),                   \
                              KvVariableShapeOp<type, ktype, vtype>);
 #define REGISTER_KERNELS_ALL(dev, type)                               \
   REGISTER_KERNELS(dev, int32, int32, type)                           \