Update vae_modules/autoencoder_wrapper.py

vae_modules/autoencoder_wrapper.py

@@ -67,7 +67,11 @@ class Autoencoder(nn.Module):
 
     def process_stable_vae(self, audio=None, embedding=None):
         if audio is not None:
+            device=audio.device
+            self.ae.encoder.to('cpu')
+            audio = audio.to('cpu')
             z = self.ae.encoder(audio)
+            z = z.to(device)
             if self.quantization_first:
                 z = self.ae.bottleneck.encode(z)
             return z