diff --git a/tests/python/quantization/test_quantization.py b/tests/python/quantization/test_quantization.py index 08303c816af1..ce43ad92fda8 100644 --- a/tests/python/quantization/test_quantization.py +++ b/tests/python/quantization/test_quantization.py @@ -510,7 +510,6 @@ def test_optimal_threshold_adversarial_case(): @with_seed() -@unittest.skip("Flaky test: https://github.com/apache/incubator-mxnet/issues/11456") def test_get_optimal_thresholds(): # Given an ndarray with elements following a uniform distribution, the optimal threshold # for quantizing the ndarray should be either abs(min(nd)) or abs(max(nd)). @@ -519,11 +518,11 @@ def get_threshold(nd): max_nd = mx.nd.max(nd) return mx.nd.maximum(mx.nd.abs(min_nd), mx.nd.abs(max_nd)).asnumpy() - nd_dict = {'layer1': mx.nd.uniform(low=-10.532, high=11.3432, shape=(8, 3, 23, 23))} + nd_dict = {'layer1': mx.nd.uniform(low=-10.532, high=11.3432, shape=(8, 3, 23, 23), dtype=np.float64)} expected_threshold = get_threshold(nd_dict['layer1']) th_dict = mx.contrib.quant._get_optimal_thresholds(nd_dict) assert 'layer1' in th_dict - assert_almost_equal(np.array([th_dict['layer1'][1]]), expected_threshold, rtol=0.001, atol=0.001) + assert_almost_equal(np.array([th_dict['layer1'][1]]), expected_threshold, rtol=1e-2, atol=1e-4) if __name__ == "__main__":