fixing map location for torch load

app.py

@@ -35,12 +35,12 @@ class SepformerFineTune(torch.nn.Module):
 
 class SourceSeparationApp:
     def __init__(self, model_path,device="cpu"):
-        self.model = self.load_model(model_path,map_location=torch.device('cpu'))
+        self.model = self.load_model(model_path)
         self.device = device
 
     def load_model(self, model_path):
         model = separator.from_hparams(source="speechbrain/sepformer-wsj03mix", savedir='pretrained_models/sepformer-wsj03mix', run_opts={"device": device})
-        checkpoint = torch.load(model_path)
+        checkpoint = torch.load(model_path, map_location=self.device)
         fine_tuned_model = SepformerFineTune(model)
         fine_tuned_model.load_state_dict(checkpoint["model"])
         return fine_tuned_model