class_docstr: Example of a TorchModelV2 using batch normalization.
code: def forward(self, input_dict, state, seq_lens):
      if isinstance(input_dict, SampleBatch):
          is_training = bool(input_dict.is_training)
      else:
          is_training = bool(input_dict.get("is_training", False))
      # Set the correct train-mode for our hidden module (only important
      # b/c we have some batch-norm layers).
      self._hidden_layers.train(mode=is_training)
      self._hidden_out = self._hidden_layers(input_dict["obs"])
      logits = self._logits(self._hidden_out)
      return logits, []
  
imports: import numpy
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.tf.misc import normc_initializer
  from ray.rllib.models.tf.tf_modelv2 import TFModelV2
  from ray.rllib.models.torch.misc import SlimFC, normc_initializer
  from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.utils.annotations import override
  from ray.rllib.utils.framework import try_import_tf, try_import_torch
label_desc: Use image node_id 2 for calling the TorchBatchNormModel obj's underlying member method code with example usage: obj.forward(input_dict, state, seq_lens) and returns: logits,  list