def compute_predictions(self, input_batch: NestedMap) -> Predictions:

    input_emb = self.softmax.emb_lookup(input_batch.ids)

    output = self.mlp_layers.fprop(input_emb)
    predictions = self.softmax.fprop(
        inputs=output,
        class_weights=input_batch.weights[:, :, jnp.newaxis],
        class_ids=input_batch.ids[:, :, jnp.newaxis])
    return predictions