Hi,
Thanks for reporting. Indeed, BeitModel currently returns an output of type BaseModelOutputWithPooling. This is a generic class that automatically generates the documentation for the model, defined here 1. However, in this case, it might be better to define a custom BeitModelOutput, that better describes the outputs of the model.
Do you mind opening a PR for this? This would mean defining a new dataclass within modeling_beit.py.
Otherwise, I’ll do it