Update model.py

model.py

@@ -17,8 +17,8 @@ class PredicateAwareSRL(nn.Module):
                  max_distance: int = 128):         # NEW: clamp |i - p| to this range for bucketing
         super().__init__()
         self.config = AutoConfig.from_pretrained(bert_name)
-        # self.bert = AutoModel.from_pretrained(bert_name)
-        self.encoder = AutoModel.from_pretrained(bert_name)
+        self.bert = AutoModel.from_pretrained(bert_name)
+        # self.encoder = AutoModel.from_pretrained(bert_name)
         self.use_indicator = use_indicator
 
         # --- Input dim to BiLSTM = BERT_dim + (indicator_dim) + (pos_dim)
@@ -76,8 +76,8 @@ class PredicateAwareSRL(nn.Module):
         device = input_ids.device
 
         # ---- BERT encoder
-        # bert_out = self.bert(
-        bert_out = self.encoder(
+        bert_out = self.bert(
+        # bert_out = self.encoder(
             input_ids=input_ids,
             token_type_ids=token_type_ids,
             attention_mask=attention_mask