fix custom op fork test (#14753)

* fix custom op fork test * trigger CI
apache · Apr 22, 2019 · 68efc15 · 68efc15
1 parent 100586a
commit 68efc15
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/tests/python/gpu/test_operator_gpu.py b/tests/python/gpu/test_operator_gpu.py
@@ -48,6 +48,7 @@
 set_default_context(mx.gpu(0))
 del test_support_vector_machine_l1_svm  # noqa
 del test_support_vector_machine_l2_svm  # noqa
+del test_custom_op_fork  #noqa
 
 
 def check_countsketch(in_dim,out_dim,n):

diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py
@@ -5393,6 +5393,9 @@ def create_operator(self, ctx, shapes, dtypes):
         x = mx.nd.Custom(length=10, depth=10, op_type="no_input_op")
     assert_almost_equal(x.asnumpy(), np.ones(shape=(10, 10), dtype=np.float32))
 
+
+@with_seed()
+def test_custom_op_fork():
     # test custom operator fork
     # see https://github.com/apache/incubator-mxnet/issues/14396
     class AdditionOP(mx.operator.CustomOp):
@@ -5430,7 +5433,7 @@ def custom_add():
         p.daemon = True
         p.start()
         p.join(5)
-        assert not p.is_alive(), "deadlock may exist in custom operator"
+        assert not p.is_alive() and p.exitcode == 0
 
 
 def _build_dot_custom(fun_forward, name):