device

app.py

@@ -12,7 +12,13 @@ from pipeline import build_audiosep
 
 CHECKPOINTS_DIR = Path("checkpoint")
 
-DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+#DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+if torch.backends.mps.is_available():
+  DEVICE = "mps"
+elif torch.cuda.is_available():
+  DEVICE = "cuda"
+else:
+  DEVICE = "cpu"
 
 # The model will be loaded in the future
 MODEL_NAME = CHECKPOINTS_DIR / "audiosep_base_4M_steps.ckpt"