Merge pull request #921 from Xilinx/hotfix/nb_cybersecurity

Fix for cybersecurity example after switch to qonnx export
Xilinx · Nov 20, 2023 · 9c40e38 · 9c40e38
2 parents ef9bb5f + 1a1e648
commit 9c40e38
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/notebooks/end2end_example/cybersecurity/1-train-mlp-with-brevitas.ipynb b/notebooks/end2end_example/cybersecurity/1-train-mlp-with-brevitas.ipynb
@@ -684,6 +684,7 @@
     "from brevitas.export import export_qonnx\n",
     "from qonnx.util.cleanup import cleanup as qonnx_cleanup\n",
     "from qonnx.core.modelwrapper import ModelWrapper\n",
+    "from qonnx.core.datatype import DataType\n",
     "from finn.transformation.qonnx.convert_qonnx_to_finn import ConvertQONNXtoFINN\n",
     "\n",
     "ready_model_filename = model_dir + \"/cybsec-mlp-ready.onnx\"\n",
@@ -708,6 +709,8 @@
     "\n",
     "# ModelWrapper\n",
     "model = ModelWrapper(ready_model_filename)\n",
+    "# Setting the input datatype explicitly because it doesn't get derived from the export function\n",
+    "model.set_tensor_datatype(model.graph.input[0].name, DataType[\"BIPOLAR\"])\n",
     "model = model.transform(ConvertQONNXtoFINN())\n",
     "model.save(ready_model_filename)\n",
     "\n",