kerasで自作の損失関数を作ったがLossがnanになる

46/46 [==============================] - 6s 56ms/step - loss: nan - accuracy: 0.1196 - val_loss: nan - val_accuracy: 0.0909
Epoch 2/10
46/46 [==============================] - 1s 26ms/step - loss: nan - accuracy: 0.1087 - val_loss: nan - val_accuracy: 0.0909
Epoch 3/10
46/46 [==============================] - 1s 27ms/step - loss: nan - accuracy: 0.1087 - val_loss: nan - val_accuracy: 0.0909
Epoch 4/10
46/46 [==============================] - 1s 26ms/step - loss: nan - accuracy: 0.1087 - val_loss: nan - val_accuracy: 0.0909

def class_balanced_cross_entropy(y_val, y_pred):
    beta=0.999
    y_val = tf.cast(y_val, tf.float32)
    y_pred = tf.cast(y_pred, tf.float32)
    y_pred = tf.clip_by_value(y_pred, 1e-6, y_pred) 

    # クラスごとのサンプル数を計算
    class_counts = tf.reduce_sum(y_val, axis=0)

    # クラスごとの重みを計算
    class_weights = 1.0 -beta / (1.0 - tf.math.pow(beta, class_counts))
    weights = class_weights /  tf.reduce_sum(class_weights) * len(class_counts)

    # 交差エントロピー
    per_class_loss = -y_val * tf.math.log(y_pred)

    # クラスごとの損失に重みをかけて合計（Class-Balenced Loss）
    class_balanced_loss = tf.reduce_sum(per_class_loss * weights,axis=1)

    # 損失の平均を計算
    mean_class_balanced_loss = tf.reduce_mean(class_balanced_loss)

    return mean_class_balanced_loss

def build_model():
  model = tf.keras.Sequential([
      layers.Input(shape=(50,95)),
      layers.Masking(mask_value=0),
      layers.LSTM(units=95,return_sequences=False),
      layers.Dense(3, activation='softmax')
  ])

  optimizer = Adam(learning_rate=0.001)

  model.compile(loss=class_balanced_cross_entropy,
                optimizer=optimizer,
                metrics=["accuracy"])
  return model

model = build_model()
history = model.fit(X_train, y_train_onehot, epochs=10, batch_size=2, validation_data=(X_test, y_test_onehot))