Am I correctly using vmap for RL-based inference?

[VladMurgoci]

    @nnx.jit
    def forward_fn(model, cards, actions):
        jax.debug.breakpoint()
        return model(cards, actions)

    action_dist_1 = jax.vmap(
        lambda action, card, player: jax.lax.cond(player == 0, lambda: forward_fn(model_1, card, action), lambda: jnp.zeros((1, 2))), in_axes=(0, 0, 0)
    )(action_history, card_tensor, env_state.current_player)

    action_dist_2 = jax.vmap(
        lambda action, card, player: jax.lax.cond(player == 1, lambda: forward_fn(model_2, card, action), lambda: jnp.zeros((1, 2))), in_axes=(0, 0, 0)
    )(action_history, card_tensor, env_state.current_player)

Jax newbie here, I'm trying to code parallel env running for a kuhn poker environment with 2 players. The problem is that envs do not share the same current player, therefore I need to figure out which model is required for each environment. Currently, I'm vmaping a function that uses a condition to figure out if each env needs an action from the 2 agents. Is this an efficient way of doing things?