kpriyanshu256 commited on
Commit
764b63c
1 Parent(s): 3a472c4

Model config fixes

Browse files
Files changed (1) hide show
  1. model.py +5 -1
model.py CHANGED
@@ -5,7 +5,11 @@ import transformers
5
  class BertAD(nn.Module):
6
  def __init__(self):
7
  super(BertAD, self).__init__()
8
- self.bert = transformers.BertModel(transformers.AutoConfig.from_pretrained('model'), output_hidden_states=True)
 
 
 
 
9
  self.layer = nn.Linear(768, 2)
10
 
11
 
 
5
  class BertAD(nn.Module):
6
  def __init__(self):
7
  super(BertAD, self).__init__()
8
+
9
+ model_config = transformers.AutoConfig.from_pretrained('model')
10
+ model_config.update({"output_hidden_states":True})
11
+
12
+ self.bert = transformers.BertModel(model_config)
13
  self.layer = nn.Linear(768, 2)
14
 
15