Add test for TopK

apache · Dec 28, 2018 · ffb5126 · ffb5126
1 parent f22df97
commit ffb5126
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/python/mxnet/contrib/onnx/onnx2mx/_op_translations.py b/python/mxnet/contrib/onnx/onnx2mx/_op_translations.py
@@ -719,4 +719,4 @@ def topk(attrs, inputs, proto_obj):
     """Returns the top k elements in an input array along the given axis."""
     new_attrs = translation_utils._add_extra_attributes(attrs,
                                                         {'ret_typ': 'both'})
-    return 'topk', new_attrs, inputs
+    return 'topk', new_attrs, inputs
diff --git a/tests/python-pytest/onnx/test_cases.py b/tests/python-pytest/onnx/test_cases.py
@@ -76,7 +76,8 @@
              'test_selu_default',
              'test_elu',
              'test_max_',
-             'test_softplus'
+             'test_softplus',
+             'test_top_k'
              ],
     'import': ['test_gather',
                'test_global_lppooling',