Update train_model.py

train_model.py

@@ -1,5 +1,5 @@
 # train_model.py (Training Script)
-
+from huggingface_hub import login
 import argparse
 from transformers import (
     GPT2Config,
@@ -145,6 +145,12 @@ def main():
         raise ValueError("HF_API_TOKEN is not set.")
 
     # Initialize tokenizer
+    try:
+        login(token=hf_token)
+        logging.info("Successfully logged in to Hugging Face Hub.")
+    except Exception as e:
+        logging.error(f"Failed to log in to Hugging Face Hub: {str(e)}")
+        raise e
     try:
         logging.info("Initializing tokenizer...")
         if args.task == "generation":