Skip to content

Commit

Permalink
[fix] Fix crash when calling unbind on evaluated tensor (#1554)
Browse files Browse the repository at this point in the history
  • Loading branch information
mfeliz-cruise authored Dec 16, 2022
1 parent 2ef6c3a commit b638e78
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 1 deletion.
2 changes: 1 addition & 1 deletion core/conversion/converters/impl/select.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ namespace impl {
namespace {

bool add_split(ConversionCtx* ctx, const torch::jit::Node* n, args& args, bool split_list, bool unbind) {
auto in = args[0].ITensor();
auto in = args[0].ITensorOrFreeze(ctx);
auto numOutputs = 1, numRemainder = 0;
std::vector<int64_t> sizes;

Expand Down
28 changes: 28 additions & 0 deletions tests/core/conversion/converters/test_select.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1122,6 +1122,34 @@ TEST(Converters, ATenUnbindNegativeAxisConvertsCorrectly) {
}
}

TEST(Converters, ATenUnbindEvaluatedTensor) {
const auto graph = R"IR(
graph(%x.1 : Tensor):
%2 : None = prim::Constant()
%3 : int[] = aten::size(%x.1)
%z.1 : Tensor = aten::zeros(%3, %2, %2, %2, %2)
%5 : int = prim::Constant[value=-1]()
%6 : Tensor[] = aten::unbind(%z.1, %5)
%o1.1 : Tensor, %o2.1 : Tensor = prim::ListUnpack(%6)
return (%o1.1, %o2.1))IR";

auto in = at::randint(1, 10, {2}, {at::kCUDA});

auto g = std::make_shared<torch::jit::Graph>();

torch::jit::parseIR(graph, g.get());

auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in});

auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in});

for (size_t i = 0; i < jit_results.size(); i++) {
auto trt = trt_results[i];
ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i].cuda(), trt, 2e-6));
}
}

TEST(Converters, ScatterValueConvertsCorrectly) {
const auto graph = R"IR(
graph(%data : Tensor,
Expand Down

0 comments on commit b638e78

Please sign in to comment.