def METHOD_NAME(self, a, s, **kwargs):
    a_pool = ops.matmul_at_b_a(s, a)
    # MinCut loss
    cut_loss = self.mincut_loss(a, s, a_pool)
    if K.ndim(a) == 3:
        cut_loss = K.mean(cut_loss)
    self.add_loss(cut_loss)
    # Post-processing of A
    a_pool = tf.linalg.set_diag(
        a_pool, tf.zeros(K.shape(a_pool)[:-1], dtype=a_pool.dtype)
    )
    a_pool = ops.normalize_A(a_pool)
    return a_pool