def compute_predictions(self, input_batch):
    """Computes predictions for `input_batch`."""
    p = self.params
    if p.model.packed_input:
      packed_input_kwargs = {
          'input_segment_ids': input_batch.src.segment_ids,
          'input_segment_pos': input_batch.src.segment_pos,
          'target_segment_ids': input_batch.tgt.segment_ids,
          'target_segment_pos': input_batch.tgt.segment_pos,
      }
    else:
      packed_input_kwargs = {}

    labels = NestedMap(
        class_ids=input_batch.tgt.labels, class_weights=input_batch.tgt.weights)
    if p.label_smoothing_prob > 0.0:
      vocab_size = p.model.softmax_tpl.num_classes
      class_probabilities = jax.nn.one_hot(labels.class_ids, vocab_size)
      fill_prob = p.label_smoothing_prob / (vocab_size - 1)
      class_probabilities = (
          (1.0 - p.label_smoothing_prob) * class_probabilities + fill_prob *
          (1.0 - class_probabilities)).astype(self.fprop_dtype)
      labels.class_probabilities = class_probabilities

    return self.model.fprop(
        inputs=input_batch.src.ids,
        input_paddings=input_batch.src.paddings,
        targets=input_batch.tgt.ids,
        target_paddings=input_batch.tgt.paddings,
        labels=labels,
        **packed_input_kwargs)