"probability/confidence" measurement of DONUT on s_rvlcdip (document classification task)

is it possible to get the “confidence”/“logit”/“probability” measure of the DONUT model on the task of document type classification task and not just the final “class/token”?
I’ve seen that the last layer of the model is Linear:

          (3): MBartDecoderLayer(
            (self_attn): MBartAttention(
              (k_proj): Linear(in_features=1024, out_features=1024, bias=True)
              (v_proj): Linear(in_features=1024, out_features=1024, bias=True)
              (q_proj): Linear(in_features=1024, out_features=1024, bias=True)
              (out_proj): Linear(in_features=1024, out_features=1024, bias=True)
            )
            (activation_fn): GELUActivation()
            (self_attn_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
            (encoder_attn): MBartAttention(
              (k_proj): Linear(in_features=1024, out_features=1024, bias=True)
              (v_proj): Linear(in_features=1024, out_features=1024, bias=True)
              (q_proj): Linear(in_features=1024, out_features=1024, bias=True)
              (out_proj): Linear(in_features=1024, out_features=1024, bias=True)
            )
            (encoder_attn_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
            (fc1): Linear(in_features=1024, out_features=4096, bias=True)
            (fc2): Linear(in_features=4096, out_features=1024, bias=True)
            (final_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
          )
        )
        (layernorm_embedding): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
        (layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
      )
    )
    (lm_head): Linear(in_features=1024, out_features=57525, bias=False)
  )
)

Has anyone found a method to do the same. If yes please share here