From 80c145a7446e40e76a2836a38ad30684e4400789 Mon Sep 17 00:00:00 2001
From: gcunhase <4861122+gcunhase@users.noreply.github.com>
Date: Wed, 26 Nov 2025 15:44:52 -0500
Subject: [PATCH] Update value info in graph

Signed-off-by: gcunhase <4861122+gcunhase@users.noreply.github.com>
---
 modelopt/onnx/autocast/precisionconverter.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/modelopt/onnx/autocast/precisionconverter.py b/modelopt/onnx/autocast/precisionconverter.py
index 38820479c..9b238f728 100644
--- a/modelopt/onnx/autocast/precisionconverter.py
+++ b/modelopt/onnx/autocast/precisionconverter.py
@@ -955,6 +955,13 @@ def _cleanup(self):
         # Remove redundant casts
         self._remove_redundant_casts()
 
+        # Update value_info in model's graph
+        self._update_value_info_in_graph()
+
+    def _update_value_info_in_graph(self):
+        for vi in self.model.graph.value_info:
+            vi.type.tensor_type.elem_type = self.value_info_map[vi.name].type.tensor_type.elem_type
+
     def _cleanup_no_consumer_nodes(self):
         network_outputs = {o.name for o in self.model.graph.output}
         nodes_to_remove = [