[LeducHoldem] suppress warnings in tests (#1202)

sotetsuk · Jul 24, 2024 · c6b9255 · c6b9255
1 parent bc9df9c
commit c6b9255
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/test_leduc_holdem.py b/tests/test_leduc_holdem.py
@@ -11,16 +11,16 @@ def make_test_state(key):
  # due to API update
  _, key = jax.random.split(key)
  rng1, rng2, rng3 = jax.random.split(key, 3)
- current_player = jnp.int8(jax.random.bernoulli(rng1))
+ current_player = jnp.int32(jax.random.bernoulli(rng1))
  init_card = jax.random.permutation(
- rng2, jnp.int8([0, 0, 1, 1, 2, 2]), independent=True
+ rng2, jnp.int32([0, 0, 1, 1, 2, 2]), independent=True
  )
  return State( # type:ignore
  _first_player=current_player,
  current_player=current_player,
  _cards=init_card[:3],
  legal_action_mask=jnp.bool_([1, 1, 0]),
- _chips=jnp.ones(2, dtype=jnp.int8),
+ _chips=jnp.ones(2, dtype=jnp.int32),
  )