From 8bf0decae24737c90e69a5234813000430e12fa2 Mon Sep 17 00:00:00 2001
From: George Ohashi <george@neuralmagic.com>
Date: Thu, 6 Mar 2025 12:20:14 -0500
Subject: [PATCH] revert train

---
 .../transformers/finetune/text_generation.py             | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/src/llmcompressor/transformers/finetune/text_generation.py b/src/llmcompressor/transformers/finetune/text_generation.py
index dadc7513b..66652b686 100644
--- a/src/llmcompressor/transformers/finetune/text_generation.py
+++ b/src/llmcompressor/transformers/finetune/text_generation.py
@@ -263,6 +263,15 @@ def main(
         # exit immediately
         return
 
+    # Training
+    if training_args.do_train:
+        checkpoint = None
+        if training_args.resume_from_checkpoint is not None:
+            checkpoint = training_args.resume_from_checkpoint
+        elif last_checkpoint is not None:
+            checkpoint = last_checkpoint
+        stage_runner.train(checkpoint)
+
     # save if model was provided as a string or custom output_dir was set
     if isinstance(model_args.model, str) or (
         training_args.output_dir