Kerasで自作の損失関数を定義するとエラーが出る。

Kerasにおいて、自作の損失関数(Spearmanの順相関係数)でモデリング構築したいのですが、学習時(estimator.fit実行時)に以下のエラーが発生してしまいます。

エラー名で検索してもダイレクトに記載されているものが自分としては見つからず、対応方法がわからず（別のサイトでは、他のセッションが動いているとこのエラーを吐くとあったので、一応OSを再起動しましたが、問題再現しました）。アドバイスをいただけないでしょうか？

ちなみに、損失関数部分は以下のリンクを見てコピペして作っています。
StackOverFLlowのリンク

エラーにreduction dimensionとあるので、途中の層の設定の不備とも考えられましたが、損失関数をデフォルトにするとエラーは消えました。

コード

Python
1# モジュールやデータ読み込みなど一部省略
2import keras.backend as K
3import tensorflow as tf
4
5# 自前損失関数の定義
6def spearman_score(y_val, y_pred):
7    return (tf.py_function(stats.spearmanr, [tf.cast(y_pred, tf.float32),
8                                       tf.cast(y_val, tf.float32)], Tout=tf.float32))
9
10def reg_model():
11  model = Sequential()
12  model.add(Dense(10, input_dim=autoscaled_X.shape[1], activation='relu'))
13  model.add(Dense(16, activation='relu'))
14  model.add(Dense(1))
15  model.summary()
16  model.compile(loss=spearman_score, optimizer='adam')
17  return model
18
19estimator = KerasRegressor(build_fn=reg_model, epochs=100, batch_size=10, verbose=0)
20estimator.fit(X_prepaired, y)

エラー内容


2021-06-09 09:24:09.395515: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:176] None of the MLIR Optimization Passes are enabled (registered 2)
2021-06-09 09:24:20.065805: W tensorflow/core/framework/op_kernel.cc:1767] OP_REQUIRES failed at reduction_ops_common.h:148 : Invalid argument: Invalid reduction dimension (1 for input with 0 dimension(s)
2021-06-09 09:24:20.065830: W tensorflow/core/framework/op_kernel.cc:1767] OP_REQUIRES failed at reduction_ops_common.h:148 : Invalid argument: Invalid reduction dimension (0 for input with 0 dimension(s)
Traceback (most recent call last):
  File "/Users/***/Documents/Projects/20210420/venv/lib/python3.7/site-packages/IPython/core/interactiveshell.py", line 3437, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-2-b0d746f88011>", line 1, in <module>
    runfile('/Users/***/Documents/Projects/20210420/modeling_with_NN.py', wdir='/Users/***/Documents/Projects/20210420')
  File "/Applications/PyCharm.app/Contents/plugins/python/helpers/pydev/_pydev_bundle/pydev_umd.py", line 197, in runfile
    pydev_imports.execfile(filename, global_vars, local_vars)  # execute the script
  File "/Applications/PyCharm.app/Contents/plugins/python/helpers/pydev/_pydev_imps/_pydev_execfile.py", line 18, in execfile
    exec(compile(contents+"\n", file, 'exec'), glob, loc)
  File "/Users/***/Documents/Projects/20210420/modeling_with_NN.py", line 120, in <module>
    estimator.fit(X_prepaired, y)
  File "/Users/***/Documents/Projects/20210420/venv/lib/python3.7/site-packages/keras/wrappers/scikit_learn.py", line 163, in fit
    history = self.model.fit(x, y, **fit_args)
  File "/Users/***/Documents/Projects/20210420/venv/lib/python3.7/site-packages/keras/engine/training.py", line 1158, in fit
    tmp_logs = self.train_function(iterator)
  File "/Users/***/Documents/Projects/20210420/venv/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 889, in __call__
    result = self._call(*args, **kwds)
  File "/Users/***/Documents/Projects/20210420/venv/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 950, in _call
    return self._stateless_fn(*args, **kwds)
  File "/Users/***/Documents/Projects/20210420/venv/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 3024, in __call__
    filtered_flat_args, captured_inputs=graph_function.captured_inputs)  # pylint: disable=protected-access
  File "/Users/***/Documents/Projects/20210420/venv/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 1961, in _call_flat
    ctx, args, cancellation_manager=cancellation_manager))
  File "/Users/***/Documents/Projects/20210420/venv/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 596, in call
    ctx=ctx)
  File "/Users/***/Documents/Projects/20210420/venv/lib/python3.7/site-packages/tensorflow/python/eager/execute.py", line 60, in quick_execute
    inputs, attrs, num_outputs)
tensorflow.python.framework.errors_impl.InvalidArgumentError:  Invalid reduction dimension (1 for input with 0 dimension(s)
	 [[node gradient_tape/sequential/dense_2/BiasAdd/Sum (defined at /Users/***/Documents/Projects/20210420/venv/lib/python3.7/site-packages/keras/engine/training.py:774) ]] [Op:__inference_train_function_600]
Function call stack:
train_function