Update preprocess_hubert_f0.py

yxys2147 · Jul 22, 2023 · 75988d0 · 75988d0
1 parent ba5d2c8
commit 75988d0
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/preprocess_hubert_f0.py b/preprocess_hubert_f0.py
@@ -32,8 +32,12 @@ def process_one(filename, hmodel,f0p,diff=False,mel_extractor=None):
     wav, sr = librosa.load(filename, sr=sampling_rate)
     audio_norm = torch.FloatTensor(wav)
     audio_norm = audio_norm.unsqueeze(0)
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
+    rank = mp.current_process()._identity
+    rank = rank[0] if len(rank) > 0 else 0
+    if torch.cuda.is_available():
+        gpu_id = rank % torch.cuda.device_count()
+        device = torch.device(f"cuda:{gpu_id}")
+    print(f"Rank {rank} uses device {device}")
     soft_path = filename + ".soft.pt"
     if not os.path.exists(soft_path):
         wav16k = librosa.resample(wav, orig_sr=sampling_rate, target_sr=16000)