Update model.py

model.py

@@ -43,7 +43,8 @@ class SALMONN(nn.Module):
         lora_dropout=0.1,
         second_per_frame=0.333333,
         second_stride=0.333333,
-        low_resource=False
+        low_resource=False,
+        device=None,
     ):
 
         super().__init__()
@@ -115,7 +116,7 @@ class SALMONN(nn.Module):
 
         # load ckpt
         ckpt_dict = torch.load(ckpt)['model']
-        self.load_state_dict(ckpt_dict, strict=False)
+        self.load_state_dict(ckpt_dict, strict=False, map_location=device)
 
     def generate(
         self,