Upload app.py

app.py

@@ -7,21 +7,22 @@ import torch
 from models import AudioClassifier
 from utils import logger
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+logger.info(f"Device: {device}")
 
 ckpt_dir = Path("ckpt/")
 config_path = ckpt_dir / "config.json"
 assert config_path.exists(), f"config.json not found in {ckpt_dir}"
 config = json.loads((ckpt_dir / "config.json").read_text())
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = AudioClassifier(device=device, **config["model"]).to(device)
 
+model = AudioClassifier(device=device, **config["model"]).to(device)
 # Latest checkpoint
 if (ckpt_dir / "model_final.pth").exists():
     ckpt = ckpt_dir / "model_final.pth"
 else:
     ckpt = sorted(ckpt_dir.glob("*.pth"))[-1]
 logger.info(f"Loading {ckpt}...")
-model.load_state_dict(torch.load(ckpt))
+model.load_state_dict(torch.load(ckpt, map_location=device))
 
 
 def classify_audio(audio_file: str):