Fix unit tests

apache · Jan 16, 2018 · 5c81473 · 5c81473
1 parent 4cf8ad7
commit 5c81473
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/python/mxnet/quantization.py b/python/mxnet/quantization.py
@@ -205,7 +205,7 @@ def _get_optimal_threshold(arr, num_bins=8001, num_quantized_bins=255):
     hist, hist_edeges = np.histogram(arr, bins=num_bins, range=(-th, th))
     zero_bin_idx = num_bins / 2
     num_half_quantized_bins = num_quantized_bins / 2
-    assert np.allclose(hist_edeges[zero_bin_idx] + hist_edeges[zero_bin_idx + 1], 0, rtol=1e-5, atol=1e-7)
+    assert np.allclose(hist_edeges[int(zero_bin_idx)] + hist_edeges[int(zero_bin_idx + 1)], 0, rtol=1e-5, atol=1e-7)
 
     thresholds = np.zeros(num_bins / 2 + 1 - num_quantized_bins / 2)
     divergence = np.zeros_like(thresholds)

diff --git a/tests/python/unittest/test_quantization.py b/tests/python/unittest/test_quantization.py
@@ -1,7 +1,9 @@
 """Some of the tests using CUDNN require a special GPU instruction called dp4a.
 Ref: http://images.nvidia.com/content/pdf/tesla/184457-Tesla-P4-Datasheet-NV-Final-Letter-Web.pdf
 """
-from mxnet.test_utils import *
+import mxnet as mx
+import numpy as np
+from mxnet.test_utils import assert_almost_equal, rand_ndarray, rand_shape_nd, same
 
 
 def test_quantize_float32_to_int8():
@@ -326,7 +328,7 @@ def test_quantize_sym_with_calib():
         assert_almost_equal(np.array([lhs]), np.array([rhs]))
         lhs = float(attr_dict[name]['max_calib_range'])
         rhs = th_dict[op_name_to_th_name[name]][1]
-        assert_almost_equal(np.array([lhs]), np.array([rhs]))
+        assert_almost_equal(np.array([lhs]), np.array([rhs]), rtol=1e-3, atol=1e-5)
 
 
 def test_get_optimal_thresholds():