class_docstr: The world-model network sub-component used to predict the `continue` flags .
  
  Predicted continue flags are used to produce "dream data" to learn the policy in.
  
  The continue flags are predicted via a linear output used to parameterize a
  Bernoulli distribution, from which simply the mode is used (no stochastic
  sampling!). In other words, if the sigmoid of the output of the linear layer is
  >0.5, we predict a continuation of the episode, otherwise we predict an episode
  terminal.
docstr: Performs a forward pass through the continue predictor.
  
  Args:
      h: The deterministic hidden state of the sequence model. [B, dim(h)].
      z: The stochastic discrete representations of the original
          observation input. [B, num_categoricals, num_classes].
docstr_tok: ["Performs","a","forward","pass","through","the","continue","predictor",".","Args",":","h",":","The","deterministic","hidden","state","of","the","sequence","model",".","[","B",",","dim","(","h",")","]",".","z",":","The","stochastic","discrete","representations","of","the","original","observation","input",".","[","B",",","num_categoricals",",","num_classes","]","."]
code: def call(self, h, z):
      """Performs a forward pass through the continue predictor.
  
      Args:
          h: The deterministic hidden state of the sequence model. [B, dim(h)].
          z: The stochastic discrete representations of the original
              observation input. [B, num_categoricals, num_classes].
      """
      # Flatten last two dims of z.
      assert len(z.shape) == 3
      z_shape = tf.shape(z)
      z = tf.reshape(z, shape=(z_shape[0], -1))
      assert len(z.shape) == 2
      out = tf.concat([h, z], axis=-1)
      out.set_shape(
          [
              None,
              (
                  get_num_z_categoricals(self.model_size)
                  * get_num_z_classes(self.model_size)
                  + get_gru_units(self.model_size)
              ),
          ]
      )
      # Send h-cat-z through MLP.
      out = self.mlp(out)
      # Remove the extra [B, 1] dimension at the end to get a proper Bernoulli
      # distribution. Otherwise, tfp will think that the batch dims are [B, 1]
      # where they should be just [B].
      logits = tf.cast(tf.squeeze(out, axis=-1), tf.float32)
      # Create the Bernoulli distribution object.
      bernoulli = tfp.distributions.Bernoulli(
          logits=logits, dtype=tf.float32
      )
  
      # Take the mode (greedy, deterministic "sample").
      continue_ = bernoulli.mode()
  
      # Return Bernoulli sample (whether to continue) OR (continue?, Bernoulli prob).
      return continue_, bernoulli
  
file_path: ray/python/ray/rllib/algorithms/dreamerv3/tf/models/components/continue_predictor.py
imports: from typing import Optional
  from ray.rllib.algorithms.dreamerv3.tf.models.components.mlp import MLP
  from ray.rllib.algorithms.dreamerv3.utils import get_gru_units, get_num_z_classes, get_num_z_categoricals
  from ray.rllib.utils.framework import try_import_tf, try_import_tfp
label_desc: Use image node_id 2 for calling the ContinuePredictor obj's underlying member method code with example usage: obj.call(h, z) and returns: continue_,  bernoulli