Skip to content

Commit 22e6a6b

Browse files
committed
feat(aten::topk): Adding BUILD files for topk op
Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
1 parent 1557f6e commit 22e6a6b

File tree

3 files changed

+7
-2
lines changed

3 files changed

+7
-2
lines changed

core/conversion/converters/impl/topk.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,8 @@ auto topk_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns().patter
2121
auto k = args[1].unwrapToInt();
2222
auto dim = args[2].unwrapToInt();
2323
auto largest = args[3].unwrapToBool();
24-
auto sorted = args[4].unwrapToBool();
25-
24+
// auto sorted = args[4].unwrapToBool(); # Currently unused
25+
2626
auto selfDim = util::toVec(self->getDimensions());
2727

2828
//reduceAxes The reduction dimensions. The bit in position i of bitmask reduceAxes corresponds to explicit dimension i of the result.

tests/core/conversion/converters/BUILD

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,10 @@ converter_test(
6767
name = "test_stack"
6868
)
6969

70+
converter_test(
71+
name = "test_topk"
72+
)
73+
7074
converter_test(
7175
name = "test_lstm_cell"
7276
)
@@ -99,5 +103,6 @@ test_suite(
99103
":test_lstm_cell",
100104
":test_unsqueeze",
101105
":test_squeeze"
106+
":test_topk",
102107
]
103108
)

0 commit comments

Comments
 (0)