You need to subclass it to have the save_pretrained methods available. So instead of
class Mean_Pooling_Model(nn.Module):
use
from transformers.modeling_utils import PreTrainedModel
class Mean_Pooling_Model(PreTrainedModel):
It will add extra functionality on top of nn.Module
.
3 Likes