fix trainer infer error

yoyolicoris · Jul 14, 2020 · d727692 · d727692
1 parent b2431e0
commit d727692
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/trainer/trainer.py b/trainer/trainer.py
@@ -56,10 +56,10 @@ def train(self):
 
                 if type(self.model) is nn.DataParallel:
                     sr = self.model.module.sr
-                    x = self.model.module.infer(mels[0], z[0].std().item())
+                    x = self.model.module.infer(mels[0], z[0].std().item())[0]
                 else:
                     sr = self.model.sr
-                    x = self.model.infer(mels[0], z[0].std().item())
+                    x = self.model.infer(mels[0], z[0].std().item())[0]
                 torch.clamp(x, -1, 1, out=x)
                 self.writer.add_audio('reconstruct_audio', x.cpu()[None, :], sample_rate=sr)