Skip to content

Commit e5b8735

Browse files
authored
Allow us to print output metadata for the case where output is not tensor
Differential Revision: D75694961 Pull Request resolved: #11263
1 parent 0e6eef8 commit e5b8735

File tree

1 file changed

+9
-1
lines changed

1 file changed

+9
-1
lines changed

extension/pybindings/pybindings.cpp

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -578,7 +578,15 @@ struct PyMethodMeta final {
578578
}
579579
py::list output_meta_strs;
580580
for (size_t i = 0; i < meta_.num_outputs(); ++i) {
581-
output_meta_strs.append(py::str(output_tensor_meta(i)->repr()));
581+
auto output_tag_res = meta_.output_tag(i);
582+
THROW_INDEX_IF_ERROR(
583+
output_tag_res.error(), "Cannot get Tag for output at %zu", i);
584+
if (output_tag_res.get() == Tag::Tensor) {
585+
output_meta_strs.append(py::str(output_tensor_meta(i)->repr()));
586+
} else {
587+
output_meta_strs.append(
588+
py::str(runtime::tag_to_string(output_tag_res.get())));
589+
}
582590
}
583591
// Add quotes to be more similar to Python's repr for strings.
584592
py::str format =

0 commit comments

Comments
 (0)