forcing torch.load to CPU

app.py

@@ -6,7 +6,8 @@ import soundfile as sf
 # forcing torch.load to CPU
 import torch
 _old_load = torch.load
-torch.load = lambda *a, **k: _old_load(*a, **{**k, 'map_location': 'cpu'})
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+torch.load = lambda *a, **k: _old_load(*a, map_location=k.pop('map_location', device))
 
 import torchaudio
 import hydra
@@ -20,7 +21,6 @@ from huggingface_hub import hf_hub_download
 import fairseq
 
 register_omegaconf_resolvers()
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 config = OmegaConf.load("configs/infer.yaml")
 
 ckpt_path = hf_hub_download(