code: def train_step_gen(
      self, training_batch: rlt.BanditRewardModelInput, batch_idx: int
  ):
      weight = self._get_sample_weight(training_batch)
  
      loss = self.reward_net.sample_elbo(
          torch.cat(
              [
                  training_batch.action,
                  training_batch.state.float_features,
              ],
              1,
          ),
          training_batch.reward,
          1,
      )
      # use (action, reward) to indicate (input,target)
  
      detached_loss = loss.detach().cpu()
      self.reporter.log(loss=detached_loss)
  
      if weight is not None:
          raise NotImplementedError  # TODO for integration in to RL framework
          # unweighted_loss = self._compute_unweighted_loss(
          #    predicted_reward, target_reward, training_batch
          # )
          # self.reporter.log(unweighted_loss=unweighted_loss)
  
      self.all_batches_processed += 1
      if self.all_batches_processed % 100 == 0:
          logger.info(
              f"{self.all_batches_processed}-th batch: "
              f"Loss={detached_loss.item()}"
          )
  
      yield loss
  
file_path: ReAgent/reagent/training/cfeval/bayes_by_backprop_trainer.py
imports: import logging
  import reagent.core.types
  import torch
  from reagent.training.cfeval.bandit_reward_network_trainer import BanditRewardNetTrainer
label_desc: Use image node_id 1 for calling the BayesByBackpropTrainer obj's underlying member method code with example usage: obj.train_step_gen(training_batch, batch_idx) without return types