hard coding indices

atomicarchitects · Aug 16, 2024 · 4606a64 · 4606a64
1 parent 6d4c641
commit 4606a64
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 7 deletions.
diff --git a/inference.cpp b/inference.cpp
@@ -13,13 +13,29 @@ int main(int argc, char *argv[]) {
     torch::inductor::AOTIModelContainerRunnerCuda *runner;
     runner = new torch::inductor::AOTIModelContainerRunnerCuda(model_path, 1);
     std::vector<torch::Tensor> inputs = {
+        // node_features
         torch::randn({32,1}, torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA)),
+
+        // pos
         torch::randn({32,3}, torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA)),
-        torch::randn({2,50}, torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA)),
-        torch::randn({32}, torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA))
+
+        // edge_index
+        torch::tensor({
+        { 1,  2,  0,  0,  3,  2,  5,  6,  4,  4,  7,  6,  9, 10,  8, 11,  8, 11,
+          9, 10, 13, 12, 14, 13, 15, 14, 17, 18, 19, 16, 16, 16, 21, 23, 20, 22,
+         21, 20, 25, 24, 26, 27, 25, 25, 29, 28, 30, 29, 31, 30},
+        { 0,  0,  1,  2,  2,  3,  4,  4,  5,  6,  6,  7,  8,  8,  9,  9, 10, 10,
+         11, 11, 12, 13, 13, 14, 14, 15, 16, 16, 16, 17, 18, 19, 20, 20, 21, 21,
+         22, 23, 24, 25, 25, 25, 26, 27, 28, 29, 29, 30, 30, 31}
+        }, torch::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA)),
+
+        // batch
+        torch::tensor({
+        0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5,
+        6, 6, 6, 6, 7, 7, 7, 7
+        }, torch::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA))
     };
     std::vector<torch::Tensor> outputs = runner->run(inputs);
-    std::cout << "Result from the first inference:"<< std::endl;
-    std::cout << outputs[0] << std::endl;
+    std::cout << "output tensor" << outputs[0] << std::endl;
     return 0;
 }
diff --git a/train.py b/train.py
@@ -100,11 +100,15 @@ def update_fn(model, opt, graphs):
                 options={"aot_inductor.output_path": os.path.join(os.getcwd(), "export/model.so"),
             })
 
-    print("Traced Shapes")
     print("node_features", graphs.numbers.shape, graphs.numbers.dtype)
     print("pos", graphs.pos.shape, graphs.pos.dtype)
-    print("edge_index", graphs.edge_index.shape, graphs.edge_index.dtype)
-    print("batch", graphs.batch.shape, graphs.batch.dtype)
+    print("edge_index", graphs.edge_index, graphs.edge_index.shape, graphs.edge_index.dtype)
+    print("batch", graphs.batch, graphs.batch.shape, graphs.batch.dtype)
+
+    runner = torch._C._aoti.AOTIModelContainerRunnerCuda(os.path.join(os.getcwd(), f"export/model.so"), 1, device)
+    outputs_export = runner.run([graphs.numbers,graphs.pos,graphs.edge_index,graphs.batch])
+    print(f"output {outputs_export[0]}")
+
 
 if __name__ == "__main__":
     train()