feat(aten::topk): Adding BUILD files for topk op

Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
pytorch · Feb 2, 2021 · 22e6a6b · 22e6a6b
1 parent 1557f6e
commit 22e6a6b
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 2 deletions.
diff --git a/core/conversion/converters/impl/topk.cpp b/core/conversion/converters/impl/topk.cpp
@@ -21,8 +21,8 @@ auto topk_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns().patter
        auto k = args[1].unwrapToInt();
        auto dim = args[2].unwrapToInt();
        auto largest = args[3].unwrapToBool();
-       auto sorted = args[4].unwrapToBool();
- 
+       // auto sorted = args[4].unwrapToBool(); # Currently unused
+
        auto selfDim = util::toVec(self->getDimensions());
 
        //reduceAxes	The reduction dimensions. The bit in position i of bitmask reduceAxes corresponds to explicit dimension i of the result.

diff --git a/tests/core/conversion/converters/BUILD b/tests/core/conversion/converters/BUILD
@@ -67,6 +67,10 @@ converter_test(
   name = "test_stack"
 )
 
+converter_test(
+  name = "test_topk"
+)
+
 converter_test(
   name = "test_lstm_cell"
 )
@@ -99,5 +103,6 @@ test_suite(
     ":test_lstm_cell",
     ":test_unsqueeze",
     ":test_squeeze"
+    ":test_topk",
   ]
 )
diff --git a/tests/core/converters/test_topk.cpp → .../core/conversion/converters/test_topk.cpp b/tests/core/converters/test_topk.cpp → .../core/conversion/converters/test_topk.cpp