add device

facebookresearch · Aug 11, 2023 · ac311af · ac311af
1 parent b948971
commit ac311af
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/audiocraft/models/multibanddiffusion.py b/audiocraft/models/multibanddiffusion.py
@@ -74,7 +74,7 @@ def get_mbd_musicgen(device=None):
         models, processors, cfgs = load_diffusion_models(path, device=device)
         DPs = []
         for i in range(len(models)):
-            schedule = NoiseSchedule(**cfgs[i].schedule, sample_processor=processors[i])
+            schedule = NoiseSchedule(**cfgs[i].schedule, sample_processor=processors[i], device=device)
             DPs.append(DiffusionProcess(model=models[i], noise_schedule=schedule))
         return MultiBandDiffusion(DPs=DPs, codec_model=codec_model)
 
@@ -106,7 +106,7 @@ def get_mbd_24khz(bw: float = 3.0, pretrained: bool = True,
         models, processors, cfgs = load_diffusion_models(path, device=device)
         DPs = []
         for i in range(len(models)):
-            schedule = NoiseSchedule(**cfgs[i].schedule, sample_processor=processors[i])
+            schedule = NoiseSchedule(**cfgs[i].schedule, sample_processor=processors[i], device=device)
             DPs.append(DiffusionProcess(model=models[i], noise_schedule=schedule))
         return MultiBandDiffusion(DPs=DPs, codec_model=codec_model)