code: def train_step_gen(
      self, training_batch: rlt.PreprocessedRankingInput, batch_idx: int
  ):
      weight = self._get_sample_weight(training_batch)
      target_reward = self._get_target_reward(training_batch)
      predicted_reward = self.reward_net(
          training_batch
      ).predicted_reward
  
      assert (
          predicted_reward.shape == target_reward.shape
          and len(target_reward.shape) == 2
          and target_reward.shape[1] == 1
      )
  
      loss = self.loss_fn(
          predicted_reward, target_reward, weight, training_batch
      )
  
      detached_loss = loss.detach().cpu()
      self.reporter.log(loss=detached_loss)
  
      if weight is not None:
          unweighted_loss = self._compute_unweighted_loss(
              predicted_reward, target_reward, training_batch
          )
          self.reporter.log(unweighted_loss=unweighted_loss)
  
      if self.all_batches_processed % 10 == 0:
          logger.info(
              f"{self.all_batches_processed}-th batch: "
              f"{self.loss_type}={detached_loss.item()}"
          )
  
      yield loss
  
imports: import logging
  from enum import Enum
  from typing import Optional
  import numpy
  import reagent.core.types
  import torch
  from reagent.core.dataclasses import field
  from reagent.models.base import ModelBase
  from reagent.optimizer.union import Optimizer__Union
  from reagent.training.reagent_lightning_module import ReAgentLightningModule
label_desc: Use image node_id 6 for calling the RewardNetTrainer obj's underlying member method code with example usage: obj.train_step_gen(training_batch, batch_idx) without return types