Update infer_tools/infer_tool.py

infer_tools/infer_tool.py

@@ -96,12 +96,12 @@ class Svc:
         @timeit
         def diff_infer():
             spk_embed = batch.get('spk_embed') if not hparams['use_spk_id'] else batch.get('spk_ids')
-            energy = batch.get('energy').cuda() if batch.get('energy') else None
+            energy = batch.get('energy').cpu() if batch.get('energy') else None
             if spk_embed is None:
-                spk_embed = torch.LongTensor([0])
+                spk_embed = torch.LongTensor([0]).cpu()
             diff_outputs = self.model(
-                hubert=batch['hubert'].cuda(), spk_embed_id=spk_embed.cuda(), mel2ph=batch['mel2ph'].cuda(),
-                f0=batch['f0'].cuda(), energy=energy, ref_mels=batch["mels"].cuda(), infer=True)
+                hubert=batch['hubert'].cpu(), spk_embed_id=spk_embed.cpu(), mel2ph=batch['mel2ph'].cpu(),
+                f0=batch['f0'].cpu(), energy=energy, ref_mels=batch["mels"].cpu(), infer=True)
             return diff_outputs
 
         outputs = diff_infer()