def METHOD_NAME(self, z, x, ce_template_mask=None, ce_keep_rate=None):
    x, aux_dict = self.forward_features(
        z,
        x,
        ce_template_mask=ce_template_mask,
        ce_keep_rate=ce_keep_rate,
    )
    return x, aux_dict