Update musc/model.py

musc/model.py

@@ -238,8 +238,9 @@ class PretrainedModel(FourHeads):
         self.download_weights(instrument)
         package_dir = os.path.dirname(os.path.realpath(__file__))
         filename = "{}_model.pt".format(instrument)
-        self.load_state_dict(torch.load(os.path.join(package_dir, filename)))
+        self.load_state_dict(torch.load(os.path.join(package_dir, filename), map_location=torch.device('cpu')))
 
+    
     def download_weights(self, instrument):
         weight_file = "{}_model.pt".format(instrument)
         package_dir = os.path.dirname(os.path.realpath(__file__))