tensorflow 47032 issue yet another reproduce https://github.com/tensorflow/tensorflow/issues/47032

tf_47032.py

import tensorflow as tf
from tensorflow import keras
import numpy as np

N = 6
CLASSES = 3
FEATURES = 4

model = keras.Sequential([
      keras.layers.Dense(10, input_shape=(FEATURES,)),
      keras.layers.Dense(CLASSES, activation='softmax')
  ])

model.compile(
  optimizer=keras.optimizers.Adam(lr=1e-3),
  loss=keras.losses.SparseCategoricalCrossentropy())

train_features = np.random.randn(N,FEATURES)
train_labels = np.random.randint(0,CLASSES,(N,1))

def pyf(x,y):
    return x,y

ds = tf.data.Dataset.from_tensor_slices((
    tf.convert_to_tensor(train_features),
    tf.convert_to_tensor(train_labels))
).map(lambda x,y: tf.py_function(
pyf, [x,y], [tf.float64, tf.int64]))


model.fit(
    ds.batch(2),    
    epochs=2,
    validation_data=ds.batch(2), 
    verbose=0,
    class_weight={c:1.0 for c in range(CLASSES)})