diff --git a/tests/python/gpu/test_operator_gpu.py b/tests/python/gpu/test_operator_gpu.py index 19fc1eca89ce..392d0cfbe3f3 100644 --- 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 index f2d8a1b2524f..cef6672519ba 100644 --- 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):