diff --git a/cbert_finetune.py b/cbert_finetune.py
index 7e1b281..bcb95d4 100755
--- a/cbert_finetune.py
+++ b/cbert_finetune.py
@@ -141,8 +141,8 @@ def main():
             _, input_ids, input_mask, segment_ids, masked_ids = batch
             """train generator at each batch"""
             optimizer.zero_grad() 
-            outputs = model(input_ids, input_mask, segment_ids,
-                    masked_lm_labels=masked_ids)
+            outputs = model(input_ids, attention_mask=input_mask, token_type_ids=segment_ids,
+                    labels=masked_ids)
             loss = outputs[0]
             loss.backward()
             avg_loss += loss.item()