diff --git a/parakeet/models/waveflow.py b/parakeet/models/waveflow.py
index c722e03..5398b22 100644
--- a/parakeet/models/waveflow.py
+++ b/parakeet/models/waveflow.py
@@ -488,7 +488,7 @@ class ConditionalWaveFlow(nn.LayerList):
         z, log_det_jacobian = self.decoder(audio, condition)
         return z, log_det_jacobian
     
-    @paddle.fluid.dygraph.no_grad
+    @paddle.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/training/experiment.py b/parakeet/training/experiment.py
index 603fbd7..8c44f93 100644
--- a/parakeet/training/experiment.py
+++ b/parakeet/training/experiment.py
@@ -75,6 +75,7 @@ class ExperimentBase(object):
         self.args = args
 
     def setup(self):
+        paddle.set_device(self.args.device)
         if self.parallel:
             self.init_parallel()