diff --git a/modules.py b/modules.py
index 34bf4c9..df34a6c 100644
--- a/modules.py
+++ b/modules.py
@@ -114,12 +114,6 @@ def mask(inputs, key_masks=None, type=None):
     array([[[ 0.0000000e+00,  0.0000000e+00, -4.2949673e+09],
         [ 0.0000000e+00,  0.0000000e+00, -4.2949673e+09]],
 
-       [[ 0.0000000e+00, -4.2949673e+09, -4.2949673e+09],
-        [ 0.0000000e+00, -4.2949673e+09, -4.2949673e+09]],
-
-       [[ 0.0000000e+00,  0.0000000e+00, -4.2949673e+09],
-        [ 0.0000000e+00,  0.0000000e+00, -4.2949673e+09]],
-
        [[ 0.0000000e+00, -4.2949673e+09, -4.2949673e+09],
         [ 0.0000000e+00, -4.2949673e+09, -4.2949673e+09]]], dtype=float32)
     """
@@ -303,4 +297,4 @@ def noam_scheme(init_lr, global_step, warmup_steps=4000.):
         until it reaches init_lr.
     '''
     step = tf.cast(global_step + 1, dtype=tf.float32)
-    return init_lr * warmup_steps ** 0.5 * tf.minimum(step * warmup_steps ** -1.5, step ** -0.5)
\ No newline at end of file
+    return init_lr * warmup_steps ** 0.5 * tf.minimum(step * warmup_steps ** -1.5, step ** -0.5)