def compute_predictions(self, input_batch: NestedMap) -> Predictions:
    """Computes predictions for `input_batch`."""
    p = self.params
    if 'tgt' in input_batch:
      input_batch = input_batch.tgt

    if 'paddings' in input_batch:
      paddings = input_batch.paddings
    else:
      paddings = jnp.equal(input_batch.segment_ids, 0).astype(self.fprop_dtype)

    if 'weights' in input_batch:
      weights = input_batch.weights
    else:
      weights = 1.0 - paddings
      weights = weights.astype(self.fprop_dtype)
      input_batch.weights = weights

    inputs = input_batch.ids
    labels = NestedMap(class_ids=input_batch.labels, class_weights=weights)
    if p.lm.packed_input:
      packed_input_kwargs = {
          'segment_ids': input_batch.segment_ids,
          'segment_pos': input_batch.segment_pos,
      }
    else:
      packed_input_kwargs = {}
    return self.lm.fprop(
        inputs=inputs,
        paddings=paddings,
        labels=labels,
        **packed_input_kwargs)