Fix dtype mismatch in embedding layer.

google · Oct 29, 2024 · d68f374 · d68f374
1 parent 56f9cdf
commit d68f374
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tf_shell_ml/embedding.py b/tf_shell_ml/embedding.py
@@ -63,9 +63,9 @@ def build(self, input_shape):
 
     def call(self, inputs):
         if inputs.dtype != tf.int64:
-            raise ValueError(
-                f"Embedding layer expects int64 input. Got {inputs.dtype}."
-            )
+            # When using model.fit() keras will cast the input to float32.
+            inputs = tf.cast(inputs, tf.int64)
+
         if inputs.ndim != 2:
             raise ValueError(f"Embedding layer expects rank 2 input. Got {inputs}.")