diff --git a/cbert_finetune.py b/cbert_finetune.py
index 7e1b281..695fc37 100755
--- a/cbert_finetune.py
+++ b/cbert_finetune.py
@@ -14,7 +14,7 @@
 import torch
 from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler
 
-from transformers import BertTokenizer, BertModel, BertForMaskedLM, AdamW, WarmupLinearSchedule
+from transformers import BertTokenizer, BertModel, BertForMaskedLM, AdamW
 #import train_text_classifier_new
 
 import cbert_utils
@@ -161,4 +161,4 @@ def main():
                 torch.save(model, save_model_path)
 
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    main()