code: def next_state_reward_dist(
      self, state: State, action: Action
  ) -> StateDistribution:
      self._model.eval()
      state_reward_tensor = (
          self._model(
              torch.tensor(state.value, dtype=torch.float)
              .reshape(-1, self._model._state_dim)
              .to(self._device),
              torch.nn.functional.one_hot(
                  torch.tensor(action.value, dtype=torch.long),
                  self._model._action_dim,
              )
              .reshape(-1, self._model._action_dim)
              .float()
              .to(self._device),
          )
          .reshape(-1)
          .cpu()
      )
      return {
          State(
              state_reward_tensor[: self._model._state_dim]
          ): RewardProbability(state_reward_tensor[-1].item())
      }
  
imports: import logging
  import gym
  import torch
  from reagent.ope.estimators.sequential_estimators import Action, ActionDistribution, ActionSpace, IPSEstimator, Model, NeuralDualDICE, RandomRLPolicy, RewardProbability, RLEstimatorInput, RLPolicy, State, StateDistribution, Transition
  from reagent.ope.utils import RunningAverage
label_desc: Use image node_id 2 for calling the ModelWrapper obj's underlying member method code with example usage: obj.next_state_reward_dist(state, action) and returns: dict