gemini-code-assist[bot] commented on code in PR #18412:
URL: https://github.com/apache/tvm/pull/18412#discussion_r2483790311
##########
docs/how_to/tutorials/export_and_load_executable.py:
##########
@@ -176,7 +176,10 @@ def forward(self, data: torch.Tensor) -> torch.Tensor: #
type: ignore[override]
# TVM returns Array objects for tuple outputs, access via indexing.
# For models imported from PyTorch, outputs are typically tuples (even for
single outputs).
# For ONNX models, outputs may be a single Tensor directly.
- result_tensor = tvm_output[0] if isinstance(tvm_output, (tuple, list))
else tvm_output
+ if hasattr(tvm_output, "__len__") and len(tvm_output) > 0:
+ result_tensor = tvm_output[0]
+ else:
+ result_tensor = tvm_output
Review Comment:
The logic to extract the result tensor is not robust. The check
`hasattr(tvm_output, "__len__")` is too broad and can lead to incorrect
behavior for single tensor outputs. For example, a `tvm.runtime.NDArray` has a
`__len__` attribute, and this check would cause it to be sliced, which is
likely not the intended behavior for a single tensor output.
A more precise way to handle this is to check if the output is a sequence
type that is not a tensor, like `tvm.runtime.Array`, `tuple`, or `list`.
I suggest the following change for a more robust implementation:
```suggestion
if isinstance(tvm_output, (tvm.runtime.Array, tuple, list)):
result_tensor = tvm_output[0]
else:
result_tensor = tvm_output
```
##########
docs/how_to/tutorials/export_and_load_executable.py:
##########
@@ -262,7 +265,10 @@ def forward(self, data: torch.Tensor) -> torch.Tensor: #
type: ignore[override]
#
# # Step 6: Extract result (output may be tuple or single Tensor)
# # PyTorch models typically return tuples, ONNX models may return a single
Tensor
-# result = output[0] if isinstance(output, (tuple, list)) else output
+# if hasattr(tvm_output, "__len__") and len(tvm_output) > 0:
+# result_tensor = tvm_output[0]
+# else:
+# result_tensor = tvm_output
Review Comment:
This example code has a few issues:
1. It uses `tvm_output`, but the variable defined in this example script is
`output` (from `output = vm["main"](input_tensor, *params)`).
2. It defines `result_tensor`, but the following print statements use
`result`. This will cause a `NameError`.
3. The logic for unpacking the output is incorrect for single tensor
outputs, similar to the issue in the main script.
I suggest correcting the variable names and using a more robust check for
the output type.
```suggestion
# if isinstance(output, (tvm.runtime.Array, tuple, list)):
# result = output[0]
# else:
# result = output
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
