From 86a500409847bd54bd3a8498545e768017106343 Mon Sep 17 00:00:00 2001
From: Damian Stewart <d@damianstewart.com>
Date: Mon, 30 Jan 2023 11:32:39 +0100
Subject: [PATCH] better main-thread detection

---
 train.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/train.py b/train.py
index 018d851..0bbcf39 100644
--- 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