better main-thread detection

2023-01-30 11:32:39 +01:00 · 2023-01-30 11:32:39 +01:00 · 86a5004098
parent 50a71e63b6
commit 86a5004098
1 changed files with 2 additions and 1 deletions
--- a/train.py
+++ b/train.py
@ -677,7 +677,8 @@ def main(args):
        """
        handles sigterm
        """
-        if threading.current_thread().__class__.__name__ == '_MainThread':
+        is_main_thread = (torch.utils.data.get_worker_info() == None)
+        if is_main_thread:
            global interrupted
            if not interrupted:
                interrupted=True