Update pico_model.py

pico_model.py

@@ -226,7 +226,8 @@ class PicoDiffusion(ClapText_Onset_2_Audio_Diffusion):
         ckpt = clap_load_state_dict(freeze_text_encoder_ckpt, skip_params=True)
         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, map_location=self.device)
+        diffusion_ckpt = torch.load(diffusion_pt)
         del diffusion_ckpt["class_emb.weight"]
         ckpt.update(diffusion_ckpt)
         self.load_state_dict(ckpt)