Merge pull request #15 from sutro-planet/fix/use-gpu

Change default pytorch device type to cuda0 默认使用GPU节点
sutro-planet · May 28, 2024 · bd4bb0d · bd4bb0d
2 parents c8f9f0a + 90ad1de
commit bd4bb0d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pytvzhen-config.json b/pytvzhen-config.json
@@ -1,5 +1,5 @@
 {
   "OUTPUT_PATH": "./output",
-  "REMOVE_BACKGROUND_MUSIC_TORCH_DEVICE": "cpu",
+  "REMOVE_BACKGROUND_MUSIC_TORCH_DEVICE": "cuda:0",
   "REMOVE_BACKGROUND_MUSIC_BASELINE_MODEL_PATH": "./models/baseline.pth"
 }