update

app/model.py

@@ -25,12 +25,14 @@ class IndoBERTBiLSTM(PreTrainedModel):
     def __init__(self, bert_config):
         super().__init__(bert_config)
         self.output_dim = OUTPUT_DIM
+        self.n_layers = 1
         self.hidden_dim = HIDDEN_DIM
         self.bidirectional = BIDIRECTIONAL
 
         self.bert = BertModel.from_pretrained(bert_path)
         self.lstm = nn.LSTM(input_size=self.bert.config.hidden_size,
                             hidden_size=self.hidden_dim,
+                            num_layers=self.n_layers,
                             bidirectional=self.bidirectional,
                             batch_first=True)
         self.dropout = nn.Dropout(DROPOUT)
@@ -39,11 +41,11 @@ class IndoBERTBiLSTM(PreTrainedModel):
     def forward(self, input_ids, attention_mask):
 
         hidden = self.init_hidden(input_ids.shape[0])
-        # print("hidden : ", type(hidden))
         output = self.bert(input_ids=input_ids, attention_mask=attention_mask)
         sequence_output = output.last_hidden_state
 
         lstm_output, (hidden_last, cn_last) = self.lstm(sequence_output, hidden)
+
         hidden_last_L=hidden_last[-2]
         hidden_last_R=hidden_last[-1]
         hidden_last_out=torch.cat([hidden_last_L,hidden_last_R],dim=-1) #[16, 1536]
@@ -72,5 +74,4 @@ class IndoBERTBiLSTM(PreTrainedModel):
                       weight.new(self.n_layers*number, batch_size, self.hidden_dim).zero_().float()
                      )
 
-        return hidden
-    
+        return hidden