Fix an error

distill_bert_to_lstm.py

@@ -147,7 +147,7 @@ def main():
     bert_model.load_state_dict(torch.load(args.bert_model_path, map_location=device))
     logger.info(f"Loaded teacher model from {args.bert_model_path}")
 
-    vocab_size = bert_model.bert.tokenizer.vocab_size
+    vocab_size = bert_train_dataset.tokenizer.vocab_size
     
     logger.info(f"LSTM Vocabulary size: {vocab_size}")
     print("LSTM Vocabulary size: ", vocab_size)