在keras模型中使用的自定义张量流梯度

Question

我有一个自定义损失函数，可以在tensorflow 1模型中很好地工作：

@tf.custom_gradient
def custom_loss(y_pred, y_true, mask):
     ... # applies mask to y_pred,y_true loss calculation and calculates z
  def grad(dy):
     ...
  return z, grad

现在，我想在顺序keras模型中使用此损失函数。

我可以在编译模型步骤中传递此函数吗？还是我也必须定义自定义图层？ 该函数采用张量，在使用keras API时我该如何处理？