Rep.distribution.choice() not working as expected

I have a bunch of prims in my stage and I want to sample one prim every step and do something with it.

I get the error:
2024-06-11 12:43:39 [Warning] [omni.graph.core.plugin] /Replicator/SDGPipeline/OgnSampleChoicePrim.outputs:samples of type rel is not compatible with /Replicator/SDGPipeline/OgnWritePrimAttribute.inputs:values of type any

import omni.replicator.core as rep

cube_path = "/World/Cube"
sphere_path = "/World/Sphere"
torus_path = "/World/Torus"

with rep.new_layer():

    with rep.trigger.on_frame():

        cube = rep.get.prim_at_path(cube_path)
        sphere = rep.get.prim_at_path(sphere_path)
        torus = rep.get.prim_at_path(torus_path)

        selected_prim = rep.distribution.choice([cube, sphere, torus])

        with rep.utils.sequential():
            with selected_prim:
                    position=rep.distribution.uniform((0, 0, 0), (2, 2, 2)))

How can I sample one prim from a list of prims in my stage?

Thank you