Update pico_model.py

pico_model.py

@@ -227,7 +227,7 @@ class PicoDiffusion(ClapText_Onset_2_Audio_Diffusion):
         del_parameter_key = ["text_branch.embeddings.position_ids"]
         ckpt = {f"freeze_text_encoder.model.{k}":v for k, v in ckpt.items() if k not in del_parameter_key}
         #diffusion_ckpt = torch.load(diffusion_pt, map_location=self.device)
-        diffusion_ckpt = torch.load(diffusion_pt)
+        diffusion_ckpt = torch.load(diffusion_pt, map_location="cpu")
         del diffusion_ckpt["class_emb.weight"]
         ckpt.update(diffusion_ckpt)
         self.load_state_dict(ckpt)