Skip to content

Instantly share code, notes, and snippets.

@khuangaf
Last active March 14, 2020 17:22
Show Gist options
  • Save khuangaf/d6e4c8b8b369d7864868884927250c56 to your computer and use it in GitHub Desktop.
Save khuangaf/d6e4c8b8b369d7864868884927250c56 to your computer and use it in GitHub Desktop.
class CustomBERTModel(BertPreTrainedModel):
def __init__(self, config, num_class):
super(CustomBERTModel, self).__init__(config)
self.bert = BertModel(config)
self.linear = nn.Linear(config.hidden_size, num_class)
model = CustomBERTModel.from_pretrained('bert-base-uncased',num_class=10)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment