fix bug

Author: fangfangssj

graph_net/torch/extractor.py

@@ -59,6 +59,8 @@ def try_rename_placeholder(node):
                 mut_graph_codes.append(gm.code)
             # 3. Generate and save model code
             base_code = gm.code
+            base_code = utils.remove_amp_pass(base_code)
+
             # gm.graph.print_tabular()
             write_code = utils.apply_templates(base_code)
             with open(os.path.join(model_path, "model.py"), "w") as fp:
@@ -90,6 +92,7 @@ def try_rename_placeholder(node):
             return gm.forward
 
         # return torch.compile(backend=extractor, dynamic=dynamic)
+        print(model)
         compiled_model = torch.compile(model, backend=extractor, dynamic=dynamic)
 
         return compiled_model

graph_net/torch/utils.py

@@ -268,3 +268,9 @@ def replay_tensor(info):
     if dtype is torch.bool:
         return (torch.randn(size=shape) > 0.5).to(dtype).to(device)
     return torch.randn(size=shape).to(dtype).to(device) * std * 0.2 + mean
+
+
+def remove_amp_pass(code):
+    lines = code.split("\n")
+    filtered_lines = [line for line in lines if "amp" not in line]
+    return "\n".join(filtered_lines)

samples/transformers-auto-model/Qwen2.5-0.5B/graph_hash.txt

@@ -1 +1 @@
-4bd91a3dcc08517dc4ac39602b20245aba8452907e56b5111bccfe53a5c1c0bf
+ad1c8cd8d22e7d12ed4a58199fd907bf24abb8b89e644651597e95a8d1102d86

samples/transformers-auto-model/Qwen2.5-0.5B/graph_net.json

@@ -1,5 +1,6 @@
 {
     "framework": "torch",
     "num_devices_required": 1,
-    "num_nodes_required": 1
+    "num_nodes_required": 1,
+    "dynamic": false
 }