hunk.code_window:     if state:
           max_seq_len = torch.max(train_batch["seq_lens"])
           mask = sequence_mask(train_batch["seq_lens"], max_seq_len)
           mask = torch.reshape(mask, [-1])
       else:
           mask = torch.ones_like(
               train_batch[Postprocessing.ADVANTAGES], dtype=torch.bool)
       policy.loss_obj = PPOLoss(
           dist_class,
           model,
           train_batch[Postprocessing.VALUE_TARGETS],
  
hunk.labels: keep keep keep keep replace replace replace keep keep keep keep keep
hunk.file_path: rllib/agents/ppo/ppo_torch_policy.py
hunk.type: replace
file: import collections
  
  from ray.util.timer import _Timer
  
  
  class MetricsContext:
      """Metrics context object for a local iterator.
  
      This object is accessible by all operators of a local iterator. It can be
      used to store and retrieve global execution metrics for the iterator.
      It can be accessed by calling LocalIterator.get_metrics(), which is only
      allowable inside iterator functions.
  
      Attributes:
          counters (defaultdict): dict storing increasing metrics.
          timers (defaultdict): dict storing latency timers.
          info (dict): dict storing misc metric values.
          current_actor (ActorHandle): reference to the actor handle that
              produced the current iterator output. This is automatically set
              for gather_async().
          parent_metrics (list): list of other MetricsContexts that have been
              attached to this due to LocalIterator.union().
      """
  
      def __init__(self):
          self.counters = collections.defaultdict(int)
          self.timers = collections.defaultdict(_Timer)
          self.info = {}
          self.current_actor = None
          self.parent_metrics = []
  
      def save(self):
          """Return a serializable copy of this context."""
          return {
              "counters": dict(self.counters),
              "info": dict(self.info),
              "timers": None,  # TODO(ekl) consider persisting timers too
              "parent_state": [u.save() for u in self.parent_metrics],
          }
  
      def restore(self, values):
          """Restores state given the output of save()."""
          self.counters.clear()
          self.counters.update(values["counters"])
          self.timers.clear()
          self.info = values["info"]
          for u, state in zip(self.parent_metrics, values["parent_state"]):
              u.restore(state)
  