入力値によって異なる損失関数を用いる方法

前提・実現したいこと

初心者です。
tensorflowを使い、入力値によって異なる損失関数を用いるディープニューラルネットワークを実装したいと思っています。
その簡単な例として入力値が0の時だけ損失関数が0となるプログラムを書いたのですが、以下のようなエラーメッセージが発生しました。
エラーの意味もよく分かっておらず、どこを直せば良いのか教えていただけるととても嬉しいです。どうぞよろしくお願い致します。

Traceback (most recent call last):
  File "/Users/XXX/.pyenv/versions/3.7.3/lib/python3.7/site-packages/tensorflow/python/ops/gen_math_ops.py", line 8631, in select_v2
    tld.op_callbacks, condition, t, e)
tensorflow.python.eager.core._FallbackException: This function does not handle the case of the path where all inputs are not already EagerTensors.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "ex11_v6.py", line 36, in <module>
    model = build_model()
  File "ex11_v6.py", line 31, in build_model
    model.compile(loss=custom_loss(train_dataset),
  File "ex11_v6.py", line 29, in custom_loss
    return tf.where(cond, zero_loss, MSE)
  File "/Users/XXX/.pyenv/versions/3.7.3/lib/python3.7/site-packages/tensorflow/python/ops/array_ops.py", line 4346, in where_v2
    return gen_math_ops.select_v2(condition=condition, t=x, e=y, name=name)
  File "/Users/XXX/.pyenv/versions/3.7.3/lib/python3.7/site-packages/tensorflow/python/ops/gen_math_ops.py", line 8636, in select_v2
    condition, t, e, name=name, ctx=_ctx)
  File "/Users/XXX/.pyenv/versions/3.7.3/lib/python3.7/site-packages/tensorflow/python/ops/gen_math_ops.py", line 8657, in select_v2_eager_fallback
    _attr_T, _inputs_T = _execute.args_to_matching_eager([t, e], ctx)
  File "/Users/XXX/.pyenv/versions/3.7.3/lib/python3.7/site-packages/tensorflow/python/eager/execute.py", line 263, in args_to_matching_eager
    t, dtype, preferred_dtype=default_dtype, ctx=ctx))
  File "/Users/XXX/.pyenv/versions/3.7.3/lib/python3.7/site-packages/tensorflow/python/framework/ops.py", line 1341, in convert_to_tensor
    ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)
  File "/Users/XXX/.pyenv/versions/3.7.3/lib/python3.7/site-packages/tensorflow/python/framework/constant_op.py", line 321, in _constant_tensor_conversion_function
    return constant(v, dtype=dtype, name=name)
  File "/Users/XXX/.pyenv/versions/3.7.3/lib/python3.7/site-packages/tensorflow/python/framework/constant_op.py", line 262, in constant
    allow_broadcast=True)
  File "/Users/XXX/.pyenv/versions/3.7.3/lib/python3.7/site-packages/tensorflow/python/framework/constant_op.py", line 270, in _constant_impl
    t = convert_to_eager_tensor(value, ctx, dtype)
  File "/Users/XXX/.pyenv/versions/3.7.3/lib/python3.7/site-packages/tensorflow/python/framework/constant_op.py", line 96, in convert_to_eager_tensor
    return ops.EagerTensor(value, ctx.device_name, dtype)
ValueError: Attempt to convert a value (<function build_model.<locals>.custom_loss.<locals>.MSE at 0x135288400>) with an unsupported type (<class 'function'>) to a Tensor.

該当のソースコード

python
1import tensorflow as tf
2import numpy as np
3
4from tensorflow import keras
5from tensorflow.keras import layers
6
7train_dataset = np.array([0.,0.,1.,2.,4.])
8train_labels = np.array([2.,1.,3.,2.,5.])
9
10def build_model():
11  model = keras.Sequential([
12    layers.Dense(2, activation='tanh', dtype='float32', input_shape=(1,)),
13    layers.Dense(1)
14  ])
15
16  def custom_loss(x):
17
18    cond = (x == 0) 
19    
20    zero_loss = tf.keras.backend.zeros((5,1))
21    
22    def MSE(y_true, y_pred):
23      error = y_true - y_pred
24      squared_loss = tf.keras.backend.square(error)
25      return squared_loss    
26    
27#    return MSE
28#    return zero_loss
29    return tf.where(cond, zero_loss, MSE)
30
31  model.compile(loss=custom_loss(train_dataset),
32                optimizer='RMSprop',
33                metrics=['mae','mse'])
34  return model
35
36model = build_model()
37
38model.summary()
39
40model.train_on_batch(train_dataset, train_labels)

試したこと

return tf.where...の行をreturn MSEとするとエラーは出なくなりますが、return zero_lossにすると以下のエラーが出ました。

Traceback (most recent call last):
  File "ex11_v6.py", line 36, in <module>
    model = build_model()
  File "ex11_v6.py", line 33, in build_model
    metrics=['mae','mse'])
  File "/Users/XXX/.pyenv/versions/3.7.3/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py", line 338, in compile
    self.loss = loss or {}  # Backwards compat.
  File "/Users/XXX/.pyenv/versions/3.7.3/lib/python3.7/site-packages/tensorflow/python/ops/resource_variable_ops.py", line 474, in __bool__
    return bool(self.read_value())
  File "/Users/XXX/.pyenv/versions/3.7.3/lib/python3.7/site-packages/tensorflow/python/framework/ops.py", line 884, in __bool__
    return bool(self._numpy())
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()