def METHOD_NAME(self, x, a, i, mask=None):
    s = self.mlp(x)
    if mask is not None:
        s *= mask[0]
    # Collapse loss
    col_loss = self.collapse_loss(a, s)
    if K.ndim(a) == 3:
        col_loss = K.mean(col_loss)
    self.add_loss(self.collapse_regularization * col_loss)
    return s