Hey @ MinishLab! 😄
Just a minor change in the distill functions where the docstring and function name mentions model_name but the signature requires a model of type PreTrainedModel. It becomes a little confusing, so would be preferable to have just model in the name and the docstring.
For example, in this function

Thanks! 😊