fix type of ev gather op

Signed-off-by: Lyaction <[email protected]>
DeepRec-AI · Apr 15, 2024 · eae2df4 · eae2df4
1 parent 524e40b
commit eae2df4
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tensorflow/python/ops/kv_variable_ops.py b/tensorflow/python/ops/kv_variable_ops.py
@@ -778,10 +778,10 @@ def sparse_read(self, indices, name=None, ev_init_value=None, counts=None):
       if self._trainable:
         tape.variable_accessed(self)
       if ev_init_value is not None:
-        default_value = ev_init_value
+        default_value = math_ops.cast(ev_init_value, self.dtype)
         is_use_default_value_tensor = True
       else:
-        default_value = ops.convert_to_tensor(1.0)
+        default_value = ops.convert_to_tensor(1.0, dtype=self.dtype)
         is_use_default_value_tensor = False
       if counts != None:
         value = gen_kv_variable_ops.kv_resource_gather_v1(self._handle,