From 96250736cc8779489ed5354e40f22ae427c7033b Mon Sep 17 00:00:00 2001
From: Anleeos <charleeos@foxmail.com>
Date: Thu, 28 Mar 2024 11:18:57 +0800
Subject: [PATCH] fix subprocess id determination

---
 data_gen/utils/process_video/extract_segment_imgs.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/data_gen/utils/process_video/extract_segment_imgs.py b/data_gen/utils/process_video/extract_segment_imgs.py
index 0a04e21..ff62d2b 100644
--- a/data_gen/utils/process_video/extract_segment_imgs.py
+++ b/data_gen/utils/process_video/extract_segment_imgs.py
@@ -308,7 +308,7 @@ def extract_segment_job(
         if "cuda" in device:
             # determine which cuda index from subprocess id
             pname = multiprocessing.current_process().name
-            pid = int(pname.rsplit("-", 1)[-1]) - 1
+            pid = 0 if pname == "MainProcess" else int(pname.rsplit("-", 1)[-1]) - 1
             cuda_id = pid % total_gpus
             device = f"cuda:{cuda_id}"