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 <[email protected]>

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