Fix model to receive input_ids
Browse files
model.py
CHANGED
@@ -52,7 +52,7 @@ class BertMesh(PreTrainedModel):
|
|
52 |
outs = torch.sigmoid(self.linear_2(outs))
|
53 |
outs = torch.flatten(outs, start_dim=1)
|
54 |
else:
|
55 |
-
cls = self.bert(input_ids=
|
56 |
outs = torch.nn.functional.relu(self.linear_1(cls))
|
57 |
outs = self.dropout_layer(outs)
|
58 |
outs = torch.sigmoid(self.linear_out(outs))
|
|
|
52 |
outs = torch.sigmoid(self.linear_2(outs))
|
53 |
outs = torch.flatten(outs, start_dim=1)
|
54 |
else:
|
55 |
+
cls = self.bert(input_ids=input_ids)[1]
|
56 |
outs = torch.nn.functional.relu(self.linear_1(cls))
|
57 |
outs = self.dropout_layer(outs)
|
58 |
outs = torch.sigmoid(self.linear_out(outs))
|