diff --git a/examples/deepvoice3/data.py b/examples/deepvoice3/data.py
index 8ab2bd37e10a5f0e0c12875fb0387063b3e1d5b7..8da3dfbec7f4ab8c0570ea0c65a18c35dc14ec75 100644
--- a/examples/deepvoice3/data.py
+++ b/examples/deepvoice3/data.py
@@ -230,7 +230,7 @@ def make_data_loader(data_root, config):
         ref_level_db=c["ref_level_db"],
         max_norm=c["max_norm"],
         clip_norm=c["clip_norm"])
-    ljspeech = CacheDataset(TransformDataset(meta, transform))
+    ljspeech = TransformDataset(meta, transform)
 
     # use meta data's text length as a sort key for the sampler
     batch_size = config["train"]["batch_size"]
diff --git a/parakeet/modules/weight_norm.py b/parakeet/modules/weight_norm.py
index 82203d6dc24eae0c096017c3566cff9c547125e6..51732a7f13fc71a6b9056f737b5dc05cda0789dd 100644
--- a/parakeet/modules/weight_norm.py
+++ b/parakeet/modules/weight_norm.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import numpy as np
 from paddle import fluid
 import paddle.fluid.dygraph as dg
 import paddle.fluid.layers as F
@@ -44,10 +43,10 @@ def norm_except(param, dim, power):
     if dim is None:
         return norm(param, dim, power)
     elif dim == 0:
-        param_matrix = F.reshape(param, (shape[0], np.prod(shape[1:])))
+        param_matrix = F.reshape(param, (shape[0], -1))
         return norm(param_matrix, dim=1, power=power)
     elif dim == -1 or dim == ndim - 1:
-        param_matrix = F.reshape(param, (np.prod(shape[:-1]), shape[-1]))
+        param_matrix = F.reshape(param, (-1, shape[-1]))
         return norm(param_matrix, dim=0, power=power)
     else:
         perm = list(range(ndim))
@@ -62,24 +61,26 @@ def compute_l2_normalized_weight(v, g, dim):
     ndim = len(shape)
 
     if dim is None:
-        v_normalized = v / (F.reduce_sum(F.square(v)) + 1e-12)
+        v_normalized = v / (F.sqrt(F.reduce_sum(F.square(v))) + 1e-12)
     elif dim == 0:
-        param_matrix = F.reshape(v, (shape[0], np.prod(shape[1:])))
+        param_matrix = F.reshape(v, (shape[0], -1))
         v_normalized = F.l2_normalize(param_matrix, axis=1)
+        v_normalized = F.reshape(v_normalized, shape)
     elif dim == -1 or dim == ndim - 1:
-        param_matrix = F.reshape(v, (np.prod(shape[:-1]), shape[-1]))
+        param_matrix = F.reshape(v, (-1, shape[-1]))
         v_normalized = F.l2_normalize(param_matrix, axis=0)
+        v_normalized = F.reshape(v_normalized, shape)
     else:
         perm = list(range(ndim))
         perm[0] = dim
         perm[dim] = 0
         transposed_param = F.transpose(v, perm)
-        param_matrix = F.reshape(
-            transposed_param,
-            (transposed_param.shape[0], np.prod(transposed_param.shape[1:])))
+        transposed_shape = transposed_param.shape
+        param_matrix = F.reshape(transposed_param,
+                                 (transposed_param.shape[0], -1))
         v_normalized = F.l2_normalize(param_matrix, axis=1)
+        v_normalized = F.reshape(v_normalized, transposed_shape)
         v_normalized = F.transpose(v_normalized, perm)
-    v_normalized = F.reshape(v_normalized, shape)
     weight = F.elementwise_mul(v_normalized, g, axis=dim)
     return weight