Skip to content

Commit 0a98180

Browse files
committed
chore: Initialize shape key as non-empty string to validate no input tensor
1 parent 377248e commit 0a98180

File tree

3 files changed

+4
-4
lines changed

3 files changed

+4
-4
lines changed

core/runtime/TRTEngine.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ struct TRTEngine : torch::CustomClassHolder {
8585
at::cuda::CUDAStream caller_stream = c10::cuda::getDefaultCUDAStream();
8686
std::vector<at::Tensor> input_buffers = {};
8787
std::vector<at::Tensor> output_buffers = {};
88-
std::string shape_key;
88+
std::string shape_key = "None";
8989
bool cudagraphs_enabled = false;
9090
bool use_pre_allocated_outputs = true;
9191
std::vector<at::Tensor> pre_allocated_outputs;

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -248,7 +248,7 @@ def create_output_tensors(self) -> List[torch.Tensor]:
248248
outputs.append(output)
249249
return outputs
250250

251-
def set_output_opt(self, enable: bool) -> None:
251+
def set_pre_allocated_outputs(self, enable: bool) -> None:
252252
self.use_pre_allocated_outputs = enable
253253

254254
def forward(self, *inputs: torch.Tensor) -> torch.Tensor | Tuple[torch.Tensor, ...]:

py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -207,7 +207,7 @@ def setup_engine(self) -> None:
207207
if self.engine is not None:
208208
return
209209
self.engine = torch.classes.tensorrt.Engine(self._pack_engine_info())
210-
self.set_output_opt(True)
210+
self.set_pre_allocated_outputs(False)
211211

212212
def encode_metadata(self, metadata: Any) -> str:
213213
metadata = copy.deepcopy(metadata)
@@ -272,7 +272,7 @@ def set_extra_state(self, state: SerializedTorchTensorRTModuleFmt) -> None:
272272
self.input_binding_names = state[2]
273273
self.output_binding_names = state[3]
274274

275-
def set_output_opt(self, enable: bool) -> None:
275+
def set_pre_allocated_outputs(self, enable: bool) -> None:
276276
self.engine.use_pre_allocated_outputs = enable
277277

278278
def forward(self, *inputs: Any) -> torch.Tensor | Tuple[torch.Tensor, ...]:

0 commit comments

Comments
 (0)