From 72e22e940a1ed19819e8b577f78ac071a207f1aa Mon Sep 17 00:00:00 2001 From: Nikil Ragav <65924881+NikilXYZ@users.noreply.github.com> Date: Mon, 13 Jan 2025 14:15:16 -0600 Subject: [PATCH] Update adaround_optimizer.py (#3709) Make sure that the input and output end up as tensors end up on the right device Signed-off-by: Nikil Ragav <65924881+NikilXYZ@users.noreply.github.com> --- .../onnx/src/python/aimet_onnx/adaround/adaround_optimizer.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/TrainingExtensions/onnx/src/python/aimet_onnx/adaround/adaround_optimizer.py b/TrainingExtensions/onnx/src/python/aimet_onnx/adaround/adaround_optimizer.py index 9603c96ffd3..835fbb6c134 100644 --- a/TrainingExtensions/onnx/src/python/aimet_onnx/adaround/adaround_optimizer.py +++ b/TrainingExtensions/onnx/src/python/aimet_onnx/adaround/adaround_optimizer.py @@ -176,6 +176,8 @@ def _optimize_rounding(cls, module: ModuleInfo, quantized_input_name, else: model_inputs = cached_dataset[np.random.randint(len(cached_dataset))] inp_data, orig_out_data = act_sampler.sample_acts(create_input_dict(orig_model.model, model_inputs)) + inp_data, orig_out_data = torch.from_numpy(inp_data[0]).to(torch_device), torch.from_numpy(out_data[0]).to(torch_device) + # This assumes there's only 1 input and 1 output in the list output by sample_acts # Clear alpha's gradients before optimization step