class_docstr: Trainer for MDN-RNN
class_docstr_tok: ["Trainer","for","MDN-RNN"]
code: def test_step(  # pyre-ignore inconsistent override because lightning doesn't use types
      self,
      training_batch: rlt.MemoryNetworkInput,
      batch_idx: int,
  ):
      (
          seq_len,
          batch_size,
          state_dim,
      ) = training_batch.state.float_features.shape
  
      losses = self.get_loss(training_batch, state_dim)
  
      detached_losses = {
          k: loss.cpu().detach().item() for k, loss in losses.items()
      }
      self.reporter.log(
          test_loss=detached_losses["loss"],
          test_gmm=detached_losses["gmm"],
          test_bce=detached_losses["bce"],
          test_mse=detached_losses["mse"],
      )
  
      loss = losses["loss"]
      self.log(
          "td_loss",
          loss,
          prog_bar=True,
          batch_size=training_batch.batch_size(),
      )
      return loss
  
file_path: ReAgent/reagent/training/world_model/mdnrnn_trainer.py
imports: import logging
  from typing import Optional
  import reagent.core.types
  import torch
  import torch.nn.functional
  from reagent.core.parameters import MDNRNNTrainerParameters
  from reagent.models.mdn_rnn import gmm_loss
  from reagent.models.world_model import MemoryNetwork
  from reagent.training.reagent_lightning_module import ReAgentLightningModule
label_desc: Use image node_id 5 for calling the MDNRNNTrainer obj's underlying member method code with example usage: obj.test_step(training_batch, batch_idx) and returns: loss