use Sort instead of TopK if K is larger than the size of the dimension we go along with

source/tensor/core/sort/TopK.cpp

@@ -23,6 +23,7 @@
 #include "../../XName.h"
 #include "../../XName.h"
 #include "TopK.h"
 #include "TopK.h"
 #include "TopK.cuh"
 #include "TopK.cuh"
+#include "Sort.h"
 
 
 namespace nts { // namespace nts(NiuTrans.Tensor)
 namespace nts { // namespace nts(NiuTrans.Tensor)
 
 
@@ -116,7 +117,10 @@ get the top-k items along a given dimension
 */
 */
 void TopK(XTensor &a, XTensor &b, XTensor &index, int dim, int k)
 void TopK(XTensor &a, XTensor &b, XTensor &index, int dim, int k)
 {
 {
-    _TopK(&a, &b, &index, dim, k);
+    if(a.dimSize[dim] <= k)
+        _Sort(&a, &b, &index, dim);
+    else
+        _TopK(&a, &b, &index, dim, k);
 
 
     /* tensor connection */
     /* tensor connection */
     XList list(2);
     XList list(2);