@@ -152,10 +152,15 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
152
152
std::vector<int32_t > inputs_cpu_vec (
153
153
input_cpu.data_ptr <int32_t >(), input_cpu.data_ptr <int32_t >() + input_cpu.numel ());
154
154
inputShapeTensorValues.emplace_back (inputs_cpu_vec);
155
- compiled_engine->exec_ctx ->setTensorAddress (name.c_str (), inputShapeTensorValues.back ().data ());
155
+ TORCHTRT_CHECK (
156
+ compiled_engine->exec_ctx ->setTensorAddress (name.c_str (), inputShapeTensorValues.back ().data ()),
157
+ " Error while setting the tensor address for shape inputs" );
156
158
} else {
157
- compiled_engine->exec_ctx ->setInputShape (name.c_str (), dims);
158
- compiled_engine->exec_ctx ->setTensorAddress (name.c_str (), inputs[i].view (shape).contiguous ().data_ptr ());
159
+ TORCHTRT_CHECK (
160
+ compiled_engine->exec_ctx ->setInputShape (name.c_str (), dims), " Error while setting the input shape" );
161
+ TORCHTRT_CHECK (
162
+ compiled_engine->exec_ctx ->setTensorAddress (name.c_str (), inputs[i].view (shape).contiguous ().data_ptr ()),
163
+ " Error while setting the input tensor address for inputs" );
159
164
}
160
165
}
161
166
@@ -188,7 +193,9 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
188
193
auto dims = core::util::toVec (out_shape);
189
194
auto type = util::TRTDataTypeToScalarType (compiled_engine->exec_ctx ->getEngine ().getTensorDataType (name.c_str ()));
190
195
outputs[pyt_idx] = std::move (at::empty (dims, {at::kCUDA }).to (type).contiguous ());
191
- compiled_engine->exec_ctx ->setTensorAddress (name.c_str (), outputs[pyt_idx].data_ptr ());
196
+ TORCHTRT_CHECK (
197
+ compiled_engine->exec_ctx ->setTensorAddress (name.c_str (), outputs[pyt_idx].data_ptr ()),
198
+ " Error while setting the output tensor address" );
192
199
}
193
200
}
194
201
0 commit comments