MyFlatten这个总是返回none导致错误

[JerryJiaGit]

看起来在应用的时候mask就是等于None的，这导致返回就是None，程序报错（model部分）。现在我直接用了kares自己的flatten，请问这个是哪里错了么？

    class MyFlatten(Layer):
        def __init__(self, **kwargs):
            self.supports_masking = True
            super(MyFlatten, self).__init__(**kwargs)

    # def compute_mask(self, inputs, mask=None):
    #     if mask==None:
    #         return mask
    # 	return k.batch_flatten(mask)
        def compute_mask(self, inputs, mask=None):
            if mask == None:
                return mask
            else:
                # return K.not_equal(inputs, 0)
                mask = k.repeat(k.not_equal(inputs, 0), self.output_dim)
                mask = tf.transpose(mask, [0, 2, 1])
                return mask
    
        def call(self, inputs, mask=None):
            return k.batch_flatten(inputs)
    
        def compute_output_shape(self, input_shape):
            return (input_shape[0], np.prod(input_shape[1:]))