Update modeling.py

modeling.py

@@ -46,7 +46,7 @@ class VoiceRestore(PreTrainedModel):
         # Optimized restoration model
         self.optimized_model = OptimizedAudioRestorationModel(device=device, bigvgan_model=self.bigvgan_model)
         save_path = "./model.safetensors"
-        state_dict = load_file(save_path, device=device)
+        state_dict = torch.load(save_path, map_location=torch.device(device))
         if 'model_state_dict' in state_dict:
             state_dict = state_dict['model_state_dict']