Update forward_pass_callback

Signed-off-by: Raj Gite <quic_rgite@quicinc.com>
quic · Dec 9, 2023 · 5dd1ba8 · 5dd1ba8
1 parent ea39024
commit 5dd1ba8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/TrainingExtensions/onnx/src/python/aimet_onnx/auto_quant_v2.py b/TrainingExtensions/onnx/src/python/aimet_onnx/auto_quant_v2.py
@@ -253,7 +253,7 @@ def __init__( # pylint: disable=too-many-arguments, too-many-locals
 
         def forward_pass_callback(session, _: Any = None):
             for input_data in tqdm(data_loader):
-                input_data_dict = utils.create_input_dict(model, input_data)
+                input_data_dict = utils.create_input_dict(model.model, input_data)
                 _ = session.run(None, input_data_dict)
 
         self.forward_pass_callback = forward_pass_callback