diff --git a/easy_rec/python/layers/backbone.py b/easy_rec/python/layers/backbone.py
index f79ce1406..e9425ed4d 100644
--- a/easy_rec/python/layers/backbone.py
+++ b/easy_rec/python/layers/backbone.py
@@ -228,8 +228,9 @@ def block_input(self, config, block_outputs, training=None, **kwargs):
         fn = eval('lambda x: x' + input_node.input_slice.strip())
         input_feature = fn(input_feature)
       if input_node.HasField('input_fn'):
-        fn = eval(input_node.input_fn)
-        input_feature = fn(input_feature)
+        with tf.name_scope(config.name):
+          fn = eval(input_node.input_fn)
+          input_feature = fn(input_feature)
       inputs.append(input_feature)
 
     if config.merge_inputs_into_list:
@@ -375,8 +376,9 @@ def call_layer(self, inputs, config, name, training, **kwargs):
           fn = eval('lambda x, i: x' + conf.input_slice.strip())
           ly_inputs = fn(ly_inputs, i)
         if conf.HasField('input_fn'):
-          fn = eval(conf.input_fn)
-          ly_inputs = fn(ly_inputs, i)
+          with tf.name_scope(config.name):
+            fn = eval(conf.input_fn)
+            ly_inputs = fn(ly_inputs, i)
         output = self.call_keras_layer(ly_inputs, name_i, training, **kwargs)
         outputs.append(output)
       if len(outputs) == 1:
diff --git a/easy_rec/python/layers/keras/fibinet.py b/easy_rec/python/layers/keras/fibinet.py
index 59ad01bd6..4a19948be 100644
--- a/easy_rec/python/layers/keras/fibinet.py
+++ b/easy_rec/python/layers/keras/fibinet.py
@@ -147,7 +147,7 @@ def build(self, input_shape):
     if field_num > 200:
       logging.warning('Too many inputs for bilinear layer: %d' % field_num)
     equal_dim = True
-    _dim = input_shape[0].shape[-1]
+    _dim = input_shape[0][-1]
     for shape in input_shape:
       assert shape.ndims == 2, 'field embeddings must be rank 2 tensors'
       if shape[-1] != _dim:
@@ -175,11 +175,11 @@ def call(self, inputs, **kwargs):
     embeddings = inputs
     field_num = len(embeddings)
 
-    # bi-linear+: p的维度为[bs, f*(f-1)/2]
+    # bi-linear+: dimension of `p` is [bs, f*(f-1)/2]
     # bi-linear:
-    # 当equal_dim=True时，p的维度为[bs, f*(f-1)/2*k]，k为embeddings的size
-    # 当equal_dim=False时，p的维度为[bs, (k_2+k_3+...+k_f)+...+(k_i+k_{i+1}+...+k_f)+...+k_f]，
-    # 其中 k_i为第i个field的embedding的size
+    #   - when equal_dim=True, dimension of `p` is [bs, f*(f-1)/2*k], k is embedding size
+    #   - when equal_dim=False, dimension of `p` is [bs, (k_2+k_3+...+k_f)+...+(k_i+k_{i+1}+...+k_f)+...+k_f],
+    #   - where k_i is the embedding size of the ith field
     if self.bilinear_type == 'all':
       v_dot = [self.dot_layer(v_i) for v_i in embeddings[:-1]]
       p = [
diff --git a/easy_rec/python/layers/keras/mask_net.py b/easy_rec/python/layers/keras/mask_net.py
index f76c9ab41..6f288f392 100644
--- a/easy_rec/python/layers/keras/mask_net.py
+++ b/easy_rec/python/layers/keras/mask_net.py
@@ -67,7 +67,7 @@ def build(self, input_shape):
         else:
           input_name = self.name
         self.input_layer_norm = lambda x: layer_norm(
-            x, name=input_name + '/input_ln', reuse=self.reuse)
+            x, name=input_name + '/input_ln', reuse=tf.AUTO_REUSE)
         self.output_layer_norm = lambda x: layer_norm(
             x, name='ln_output', reuse=self.reuse)
     if self.config.HasField('output_size'):
diff --git a/easy_rec/python/test/train_eval_test.py b/easy_rec/python/test/train_eval_test.py
index dd2bd8f38..a4e1e6a1f 100644
--- a/easy_rec/python/test/train_eval_test.py
+++ b/easy_rec/python/test/train_eval_test.py
@@ -883,7 +883,7 @@ def test_train_with_multi_worker_mirror(self):
   @unittest.skipIf(
       LooseVersion(tf.__version__) != LooseVersion('2.3.0'),
       'MultiWorkerMirroredStrategy need tf version == 2.3')
-  def test_train_with_multi_worker_mirror(self):
+  def test_train_mmoe_with_multi_worker_mirror(self):
     self._success = test_utils.test_distributed_train_eval(
         'samples/model_config/mmoe_mirrored_strategy_on_taobao.config',
         self._test_dir)