def _find_one_after_last_labeled_index(self):
        body_index = 0
        while self.bodies[body_index].segment_id in self.state.body_labels:
            body_index += 1
        return body_index
def sample_fn(x):
    return utils.optimise_and_sample(x, module=model,
                                     data=None, is_training=False)[0]