Skip to content

Commit eab71f3

Browse files
committed
fix: Minor bugfix in partitioning test
- Partitioning test incorrectly expected 1 conditional engine, but got 2 since `log_sigmoid` operator is not currently supported
1 parent a32e254 commit eab71f3

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

tests/core/partitioning/test_conditionals.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,12 +35,13 @@ TEST(Partitioning, FallbackOnConditionalsCorrectly) {
3535
auto g = mod.get_method("forward").graph();
3636
torch_tensorrt::core::CompileSpec cfg(inputs);
3737
cfg.partitioning_info.enabled = true;
38+
cfg.partitioning_info.forced_fallback_operators.push_back("aten::log_sigmoid");
3839
torch::jit::script::Module new_mod = torch_tensorrt::core::CompileGraph(mod, cfg);
3940
auto new_g = new_mod.get_method("forward").graph();
4041

4142
auto conditional_engines_count = count_trt_engines_in_conditionals(new_g);
4243

43-
ASSERT_TRUE(conditional_engines_count == 1);
44+
ASSERT_TRUE(conditional_engines_count == 2);
4445
}
4546

4647
TEST(Partitioning, FallbackInplaceOPInConditionalsCorrectly) {

0 commit comments

Comments
 (0)