From 3c38ba4415b4af3d2c4fc624c96625c5ea73b86f Mon Sep 17 00:00:00 2001
From: Jack Zhou <zhoushunjie@baidu.com>
Date: Tue, 27 Dec 2022 16:02:57 +0800
Subject: [PATCH] [FDTensor] Add device id to output tensor (#978)

Add device id
---
 fastdeploy/runtime.cc | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/fastdeploy/runtime.cc b/fastdeploy/runtime.cc
index c5521f1ad..0110b4313 100755
--- a/fastdeploy/runtime.cc
+++ b/fastdeploy/runtime.cc
@@ -660,7 +660,11 @@ bool Runtime::Infer(std::vector<FDTensor>& input_tensors,
 }
 
 bool Runtime::Infer() {
-  return backend_->Infer(input_tensors_, &output_tensors_, false);
+  bool result = backend_->Infer(input_tensors_, &output_tensors_, false);
+  for (auto& tensor : output_tensors_) {
+    tensor.device_id = option.device_id;
+  }
+  return result;
 }
 
 void Runtime::BindInputTensor(const std::string& name, FDTensor& input) {