update feature check

checkpoint.py

@@ -7,9 +7,7 @@ from transformers import (
     TrainingArguments, 
     Trainer, 
     DataCollatorForLanguageModeling, 
-    TrainerCallback, 
-    TrainerState, 
-    TrainerControl
+    TrainerCallback
 )
 from peft import LoraConfig, get_peft_model
 import spaces
@@ -83,8 +81,9 @@ class SaveCheckpointCallback(TrainerCallback):
         if state.global_step % args.save_steps == 0 and state.global_step != 0:
             checkpoint_path = os.path.join(args.output_dir, f"checkpoint-{state.global_step}")
             print(f"Lưu checkpoint tại: {checkpoint_path}")
+            trainer = kwargs['trainer']  # Truy cập trainer từ kwargs
             trainer.save_model(checkpoint_path)
-        return TrainerControl.CONTINUE
+        return control  # Trả về đối tượng control hiện tại
 
 # Định Nghĩa Hàm Huấn Luyện với Decorator @spaces.GPU
 @spaces.GPU(duration=15, queue=False)