code: def get_step_prediction(
      step_predict_network: FullyConnectedNetwork,
      training_batch: rlt.MemoryNetworkInput,
  ):
      first_step_state = training_batch.state.float_features[0]
      pred_step = step_predict_network(first_step_state)
      step_probability = F.softmax(pred_step, dim=1)
      return step_probability
  
imports: import logging
  import reagent.core.types
  import torch
  import torch.nn
  import torch.nn.functional
  from reagent.core.parameters import Seq2RewardTrainerParameters
  from reagent.models.fully_connected_network import FullyConnectedNetwork
  from reagent.models.seq2reward_model import Seq2RewardNetwork
  from reagent.training.reagent_lightning_module import ReAgentLightningModule
  from reagent.training.utils import gen_permutations
label_desc: Use image node_id 1 for calling a global function with example usage: get_step_prediction(step_predict_network, training_batch) and returns: step_probability