Keras 自定义 loss 的维度和 sample weight 维度不匹配问题

在使用 LSTM＋CRF 做 NER 时，使用的 CRF 实现来自简明条件随机场CRF介绍，loss 的输出维度是 (bs,) 而由于自己不清楚 mask 的使用方式，在 Embedding 层中加入了 mask_zero=True，同时在后面使用了 Masking 层，导致出现以下错误：

tensorflow.python.framework.errors_impl.InvalidArgumentError:  assertion failed: [weights can not be broadcast to values.] [weights.shape=] [loss/crf_loss/Cast:0] [128 52] [values.shape=] [loss/crf_loss/sub_1:0] [128 1] [is_scalar=] [0]
	 [[{{node loss/crf_loss/weighted_loss/broadcast_weights/assert_broadcastable/AssertGuard/else/_164/Assert}}]] [Op:__inference_distributed_function_19323]

Function call stack:
distributed_function

sample weight 的维度应该是由 mask 造成，去掉 mask 这个错误消失

tensorflow 版本：2.1