Skip to content

Commit

Permalink
fix: Minor bugfix in partitioning test
Browse files Browse the repository at this point in the history
- Partitioning test incorrectly expected 1 conditional engine, but got 2
since `log_sigmoid` operator is not currently supported
  • Loading branch information
gs-olive committed Feb 22, 2023
1 parent a32e254 commit eab71f3
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion tests/core/partitioning/test_conditionals.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -35,12 +35,13 @@ TEST(Partitioning, FallbackOnConditionalsCorrectly) {
auto g = mod.get_method("forward").graph();
torch_tensorrt::core::CompileSpec cfg(inputs);
cfg.partitioning_info.enabled = true;
cfg.partitioning_info.forced_fallback_operators.push_back("aten::log_sigmoid");
torch::jit::script::Module new_mod = torch_tensorrt::core::CompileGraph(mod, cfg);
auto new_g = new_mod.get_method("forward").graph();

auto conditional_engines_count = count_trt_engines_in_conditionals(new_g);

ASSERT_TRUE(conditional_engines_count == 1);
ASSERT_TRUE(conditional_engines_count == 2);
}

TEST(Partitioning, FallbackInplaceOPInConditionalsCorrectly) {
Expand Down

0 comments on commit eab71f3

Please sign in to comment.