Update model_split.py

model_split.py

@@ -60,6 +60,6 @@ def auto_split_model(model, num_stages=3, output_dir="model_stage_files"):
 
 # === Example Usage ===
 if __name__ == "__main__":
-    #from transformers import GPT2Model
-    #model = GPT2Model.from_pretrained("gpt2")
+    from transformers import GPT2Model
+    model = GPT2Model.from_pretrained("gpt2")
     auto_split_model(model, num_stages=3)