diff --git a/paddleslim/prune/criterion.py b/paddleslim/prune/criterion.py
index 36410ee8d9f242c7a5f8f6d11977b3291212da02..5f1f78c1baaa4740036d7bf1780ee29dbc95789a 100644
--- a/paddleslim/prune/criterion.py
+++ b/paddleslim/prune/criterion.py
@@ -17,7 +17,7 @@
 import logging
 import numpy as np
 from ..common import get_logger
-from ..core import Registry
+from ..core import Registry, GraphWrapper
 
 __all__ = ["l1_norm", "CRITERION"]
 
@@ -61,13 +61,16 @@ def geometry_median(group, graph):
     channel_num = value.shape[0]
     w.shape = value.shape[0], np.product(value.shape[1:])
     x = w.repeat(channel_num, axis=0)
-    y = np.tile(channel_num, (channel_num, 1))
+    y = np.zeros_like(x)
+    for i in range(channel_num):
+        y[i * channel_num:(i + 1) * channel_num] = np.tile(channel_num,
+                                                           (channel_num, 1))
     tmp = np.sqrt(np.sum((x - y)**2, -1))
     tmp = tmp.reshape((channel_num, channel_num))
     tmp = np.sum(tmp, -1)
 
     for name, value, axis in group:
-        scores.append(name, axis, tmp)
+        scores.append((name, axis, tmp))
     return scores
 
 
diff --git a/paddleslim/prune/idx_selector.py b/paddleslim/prune/idx_selector.py
index 58cf1111aad35bc5101a6ab5a670c7ef1235360f..a46a76ab1ba108604ef7a8403a35c50ef3a65350 100644
--- a/paddleslim/prune/idx_selector.py
+++ b/paddleslim/prune/idx_selector.py
@@ -96,7 +96,7 @@ def optimal_threshold(group, ratio):
     name, axis, score = group[
         0]  # sort channels by the first convolution's score
 
-    score[scoew < 1e-18] = 1e-18
+    score[score < 1e-18] = 1e-18
     score_sorted = np.sort(score)
     score_square = score_sorted**2
     total_sum = score_square.sum()
diff --git a/tests/test_optimal_threshold.py b/tests/test_optimal_threshold.py
index 4f01b27082820a02a4a3ccd3c53b86c46ef84444..ca9834e355b3f66c9ddead4f0b7edf9857375337 100644
--- a/tests/test_optimal_threshold.py
+++ b/tests/test_optimal_threshold.py
@@ -49,7 +49,8 @@ class TestPrune(unittest.TestCase):
         exe = fluid.Executor(place)
         scope = fluid.Scope()
         exe.run(startup_program, scope=scope)
-        criterion = 'optimal_threshold'
+        criterion = 'bn_scale'
+        idx_selector = 'optimal_threshold'
         pruner = Pruner(criterion)
         main_program, _, _ = pruner.prune(
             main_program,
diff --git a/tests/test_prune.py b/tests/test_prune.py
index 60fe603ccd04ec1fbda9d37edc117d0b93c8b1fd..3d8bc50cf0d2d9b1959a3b9314c8e94a00ca9d06 100644
--- a/tests/test_prune.py
+++ b/tests/test_prune.py
@@ -15,7 +15,7 @@ import sys
 sys.path.append("../")
 import unittest
 import paddle.fluid as fluid
-from paddleslim.prune.walk_pruner import Pruner
+from paddleslim.prune import Pruner
 from layers import conv_bn_layer
 
 
@@ -72,7 +72,8 @@ class TestPrune(unittest.TestCase):
 
         for param in main_program.global_block().all_parameters():
             if "weights" in param.name:
-                print("param: {}; param shape: {}".format(param.name, param.shape))
+                print("param: {}; param shape: {}".format(param.name,
+                                                          param.shape))
                 self.assertTrue(param.shape == shapes[param.name])
 
 
diff --git a/tests/test_slim_prune.py b/tests/test_slim_prune.py
index 943c1c94de0ddf47f4f2a0aec40918b8f0876850..15690619e7f62eccbb50730da40a7dd7573ef081 100644
--- a/tests/test_slim_prune.py
+++ b/tests/test_slim_prune.py
@@ -49,7 +49,7 @@ class TestPrune(unittest.TestCase):
         exe = fluid.Executor(place)
         scope = fluid.Scope()
         exe.run(startup_program, scope=scope)
-        criterion = 'batch_norm_scale'
+        criterion = 'bn_scale'
         pruner = Pruner(criterion)
         main_program, _, _ = pruner.prune(
             main_program,