This is literally the Relay code Torch frontend generates. Is this what you are asking? Or do you want a Python code written by hand that is equivalent to the program below?
fn (%states: List[Tensor[(2, 4), float32]]) -> Tensor[(2, 4), float32] {
@nth(%states, 0 /* ty=int32 */) /* ty=Tensor[(2, 4), float32] */
}