From 382681a35d41df9d9d8e83c97d64415dc0ec58ea Mon Sep 17 00:00:00 2001 From: James Betker Date: Fri, 4 Mar 2022 13:42:14 -0700 Subject: [PATCH] Load diffusion_fid DVAE into the correct cuda device --- codes/trainer/eval/audio_diffusion_fid.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codes/trainer/eval/audio_diffusion_fid.py b/codes/trainer/eval/audio_diffusion_fid.py index 62e4ccbc..89c5059c 100644 --- a/codes/trainer/eval/audio_diffusion_fid.py +++ b/codes/trainer/eval/audio_diffusion_fid.py @@ -48,7 +48,7 @@ class AudioDiffusionFid(evaluator.Evaluator): if mode == 'tts': self.diffusion_fn = self.perform_diffusion_tts elif mode == 'vocoder': - self.dvae = load_speech_dvae() + self.dvae = load_speech_dvae().to(self.env['device']) self.dvae.eval() self.diffusion_fn = self.perform_diffusion_vocoder