diff --git a/parakeet/models/waveflow.py b/parakeet/models/waveflow.py
index 1274c47..89bdbda 100644
--- a/parakeet/models/waveflow.py
+++ b/parakeet/models/waveflow.py
@@ -483,7 +483,7 @@ class ConditionalWaveFlow(nn.LayerList):
         z, log_det_jacobian = self.decoder(audio, condition)
         return z, log_det_jacobian
     
-    @paddle.no_grad()
+    @paddle.fluid.dygraph.no_grad
     def synthesize(self, mel):
         condition = self.encoder(mel, trim_conv_artifact=True) #(B, C, T)
         batch_size, _, time_steps = condition.shape
diff --git a/parakeet/modules/losses.py b/parakeet/modules/losses.py
index ad29e0d..9dd40f0 100644
--- a/parakeet/modules/losses.py
+++ b/parakeet/modules/losses.py
@@ -23,9 +23,3 @@ def masked_l1_loss(prediction, target, mask):
 def masked_softmax_with_cross_entropy(logits, label, mask, axis=-1):
     ce = F.softmax_with_cross_entropy(logits, label, axis=axis)
     return weighted_mean(ce, mask)
-
-
-
-
-    
-