Fix dtype of tf argmax

keras-team · Nov 30, 2024 · 7943660 · 7943660
1 parent f529d62
commit 7943660
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/keras/src/trainers/data_adapters/tf_dataset_adapter.py b/keras/src/trainers/data_adapters/tf_dataset_adapter.py
@@ -128,7 +128,7 @@ def class_weights_map_fn(*data):
         if y.shape.rank >= 2:
             y_classes = tf.__internal__.smart_cond.smart_cond(
                 tf.shape(y)[-1] > 1,
-                lambda: tf.argmax(y, axis=-1),
+                lambda: tf.argmax(y, axis=-1, output_type=tf.int32),
                 lambda: tf.cast(tf.round(tf.squeeze(y, axis=-1)), tf.int32),
             )
         else: