From fd7ce739752aa01b248813ba5f244915940ffff6 Mon Sep 17 00:00:00 2001
From: Giuseppe Franco <giuseppefranco4@gmail.com>
Date: Tue, 28 Nov 2023 17:27:36 +0000
Subject: [PATCH] Fix (export): shape propagation in debug flow

---
 src/brevitas/export/onnx/debug.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/brevitas/export/onnx/debug.py b/src/brevitas/export/onnx/debug.py
index 85acba31b..3d5c236ac 100644
--- a/src/brevitas/export/onnx/debug.py
+++ b/src/brevitas/export/onnx/debug.py
@@ -14,6 +14,7 @@ class DebugMarkerFunction(Function):
     @staticmethod
     def symbolic(g, input, export_debug_name):
         ret = g.op('brevitas.onnx::DebugMarker', input, export_debug_name_s=export_debug_name)
+        ret.setType(input.type())
         return ret
 
     @staticmethod