poke-env: Error encountered during player.ladder()

Hi, I was testing a model I trained on Pokemon Showdown (code snippet below) when I ran into this issue. I’m able to challenge the bot to a battle and play against it perfectly well but when I do player.ladder(100) it errors out after completing a single battle.

2022-07-25 18:33:47,574 - UABGLSimpleDQN - ERROR - Unhandled exception raised while handling message:
>battle-gen8randombattle-1625188644
|-message|Nukkumatti lost due to inactivity.
|
|win|UABGLSimpleDQN
Traceback (most recent call last):
  File "E:\Dev\meta-discovery\torch_env\lib\site-packages\poke_env\player\player_network_interface.py", line 131, in _handle_message
    await self._handle_battle_message(split_messages)
  File "E:\Dev\meta-discovery\torch_env\lib\site-packages\poke_env\player\player.py", line 235, in _handle_battle_message
    self._battle_finished_callback(battle)
  File "E:\Dev\meta-discovery\torch_env\lib\site-packages\poke_env\player\env_player.py", line 106, in _battle_finished_callback
    self._observations[battle].put(self.embed_battle(battle))
KeyError: <poke_env.environment.battle.Gen8Battle object at 0x000001E1988D2EA0>
Task exception was never retrieved
future: <Task finished name='Task-39' coro=<PlayerNetwork._handle_message() done, defined at E:\Dev\meta-discovery\torch_env\lib\site-packages\poke_env\player\player_network_interface.py:117> exception=KeyError(<poke_env.environment.battle.Gen8Battle object at 0x000001E1988D2EA0>)>
Traceback (most recent call last):
  File "E:\Dev\meta-discovery\torch_env\lib\site-packages\poke_env\player\player_network_interface.py", line 177, in _handle_message
    raise exception
  File "E:\Dev\meta-discovery\torch_env\lib\site-packages\poke_env\player\player_network_interface.py", line 131, in _handle_message
    await self._handle_battle_message(split_messages)
  File "E:\Dev\meta-discovery\torch_env\lib\site-packages\poke_env\player\player.py", line 235, in _handle_battle_message
    self._battle_finished_callback(battle)
  File "E:\Dev\meta-discovery\torch_env\lib\site-packages\poke_env\player\env_player.py", line 106, in _battle_finished_callback
    self._observations[battle].put(self.embed_battle(battle))
KeyError: <poke_env.environment.battle.Gen8Battle object at 0x000001E1988D2EA0>

Model code:

class SimpleRLPlayerTesting(SimpleRLPlayer):
    def __init__(self, model, *args, **kwargs):
        SimpleRLPlayer.__init__(self, *args, **kwargs)
        self.model = model

    def choose_move(self, battle):
        state = self.embed_battle(battle)
        with torch.no_grad():
            predictions = self.model(state)
        action_mask = self.action_masks()
        action = np.argmax(predictions + action_mask)
        return self._action_to_move(action, battle)

Script:

async def main():
    ...
    player = simple_agent.SimpleRLPlayerTesting(
        model=model,
        player_configuration=PlayerConfiguration(USERNAME, PASSWORD),
        server_configuration=ShowdownServerConfiguration,
        start_timer_on_battle_start=True,
        **player_kwargs
    )
    print("Connecting to Pokemon Showdown...")
    await player.ladder(NUM_GAMES)
    # Print the rating of the player and its opponent after each battle
    for battle in player.battles.values():
        print(battle.rating, battle.opponent_rating)

if __name__ == "__main__":
    asyncio.get_event_loop().run_until_complete(main())