def modify_state_for_body(s, body):
            s.layers['segmentation'].segments = frozenset([body.segment_id])
            s.voxel_coordinates = body.bbox_start + body.bbox_size // 2
def sample_fn(x):
    return utils.optimise_and_sample(x, module=model,
                                     data=None, is_training=False)[0]