ing0 commited on
Commit
ee20368
·
1 Parent(s): 5e4e352
Files changed (1) hide show
  1. diffrhythm/infer/infer_utils.py +1 -0
diffrhythm/infer/infer_utils.py CHANGED
@@ -34,6 +34,7 @@ def prepare_model(device):
34
 
35
  # prepare vae
36
  vae_ckpt_path = hf_hub_download(repo_id="ASLP-lab/DiffRhythm-vae", filename="vae_model.pt")
 
37
  vae = torch.jit.load(vae_ckpt_path).to(device)
38
 
39
  return cfm, tokenizer, muq, vae
 
34
 
35
  # prepare vae
36
  vae_ckpt_path = hf_hub_download(repo_id="ASLP-lab/DiffRhythm-vae", filename="vae_model.pt")
37
+ print(f"****************** {device} ******************")
38
  vae = torch.jit.load(vae_ckpt_path).to(device)
39
 
40
  return cfm, tokenizer, muq, vae