diff --git a/src/plugins/proxy/src/infer_request.cpp b/src/plugins/proxy/src/infer_request.cpp index 308631a6b59..7babcadb546 100644 --- a/src/plugins/proxy/src/infer_request.cpp +++ b/src/plugins/proxy/src/infer_request.cpp @@ -51,7 +51,13 @@ ov::Tensor ov::proxy::InferRequest::get_tensor(const ov::Output& } void ov::proxy::InferRequest::set_tensor(const ov::Output& port, const ov::Tensor& tensor) { - m_infer_request->set_tensor(port, tensor); + auto new_tensor = tensor; + if (tensor.is()) { + auto remote_context = std::dynamic_pointer_cast(m_compiled_model->get_context()); + OPENVINO_ASSERT(remote_context); + new_tensor = remote_context->wrap_tensor(tensor.as()); + } + m_infer_request->set_tensor(port, new_tensor); } std::vector ov::proxy::InferRequest::get_tensors(const ov::Output& port) const { @@ -68,7 +74,15 @@ std::vector ov::proxy::InferRequest::get_tensors(const ov::Output& port, const std::vector& tensors) { - return m_infer_request->set_tensors(port, tensors); + auto new_tensors = tensors; + for (auto&& tensor : new_tensors) { + if (tensor.is()) { + auto remote_context = std::dynamic_pointer_cast(m_compiled_model->get_context()); + OPENVINO_ASSERT(remote_context); + tensor = remote_context->wrap_tensor(tensor.as()); + } + } + return m_infer_request->set_tensors(port, new_tensors); } std::vector> ov::proxy::InferRequest::query_state() const {