Tom Aarsen commited on
Commit
12e66dc
·
1 Parent(s): dd59cdc

Add base_model_prefix="model" on PreTrainedModel

Browse files

This allows you to load the AutoModel, even though the model checkpoint is for the Masked Language modeling head. In short: the base_model_prefix is what it searches for in the subclasses like NeoBERTLMHead and NeoBERTForSequenceClassification

Files changed (1) hide show
  1. model.py +1 -0
model.py CHANGED
@@ -199,6 +199,7 @@ class EncoderBlock(nn.Module):
199
 
200
  class NeoBERTPreTrainedModel(PreTrainedModel):
201
  config_class = NeoBERTConfig
 
202
  _supports_cache_class = True
203
 
204
  def _init_weights(self, module):
 
199
 
200
  class NeoBERTPreTrainedModel(PreTrainedModel):
201
  config_class = NeoBERTConfig
202
+ base_model_prefix = "model"
203
  _supports_cache_class = True
204
 
205
  def _init_weights(self, module):