No Cuda Stuff for Zero

models/salmonn.py

@@ -66,7 +66,7 @@ class SALMONN(nn.Module):
 
         # beats
         self.beats_ckpt = beats_path
-        beats_checkpoint = torch.load(self.beats_ckpt, map_location=device)
+        beats_checkpoint = torch.load(self.beats_ckpt)
         beats_cfg = BEATsConfig(beats_checkpoint["cfg"])
         beats = BEATs(beats_cfg)
         beats.load_state_dict(beats_checkpoint["model"])
@@ -130,7 +130,7 @@ class SALMONN(nn.Module):
         ).to(device)
 
         # load ckpt
-        ckpt_dict = torch.load(ckpt)["model"]
+        ckpt_dict = torch.load(ckpt, map_location="cpu")["model"]
         self.load_state_dict(ckpt_dict, strict=False)
 
     def generate(