From 08c81ad4ab7429297e14b02353f5ddddb876aa56 Mon Sep 17 00:00:00 2001
From: lixy9474 <lxy268263@alibaba-inc.com>
Date: Fri, 17 Feb 2023 09:20:22 +0800
Subject: [PATCH] [Embedding] Fix bug of saving EmbeddingVariable with int32
 type. (#692)

---
 tensorflow/core/kernels/save_restore_v2_ops.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/save_restore_v2_ops.cc b/tensorflow/core/kernels/save_restore_v2_ops.cc
index b2c651b03b8..2dbb05bad13 100644
--- a/tensorflow/core/kernels/save_restore_v2_ops.cc
+++ b/tensorflow/core/kernels/save_restore_v2_ops.cc
@@ -171,7 +171,8 @@ class SaveV2 : public OpKernel {
       const string& tensor_name = tensor_names_flat(i);
       if (tensor_types_[i] == DT_RESOURCE) {
         auto& handle = HandleFromInput(context, i + kFixedInputs);
-        if (IsHandle<EmbeddingVar<int64, float>>(handle)) {
+        if (IsHandle<EmbeddingVar<int64, float>>(handle) ||
+            IsHandle<EmbeddingVar<int32, float>>(handle)) {
           if (ev_key_types_[start_ev_key_index] == DT_INT32) {
             DumpEvWithGlobalStep<int32, float>(context,
                 i + kFixedInputs, tensor_name, writer, tensor_types_[0]);