Fix ONNX AdaRound Actn Sampler Call (#2626)

Signed-off-by: Raj Gite <quic_rgite@quicinc.com>
quic · Dec 28, 2023 · 6b66b20 · 6b66b20
1 parent 48e8781
commit 6b66b20
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/TrainingExtensions/onnx/src/python/aimet_onnx/adaround/adaround_optimizer.py b/TrainingExtensions/onnx/src/python/aimet_onnx/adaround/adaround_optimizer.py
@@ -173,7 +173,7 @@ def _optimize_rounding(cls, module: ModuleInfo, quantized_input_name,
                 orig_out_data = all_orig_out_data[indices].to(device)
             else:
                 model_inputs = cached_dataset[np.random.randint(len(cached_dataset))]
-                inp_data, orig_out_data = act_sampler.sample_acts(model_inputs)
+                inp_data, orig_out_data = act_sampler.sample_acts(create_input_dict(orig_model.model, model_inputs))
 
 
             # Clear alpha's gradients before optimization step