fixed mem monitor device (#433)

fixed mem monitor device
2022-03-16 15:25:02 +08:00 · 2022-03-16 15:25:02 +08:00 · b03b3ae99c
parent 14a7094243
commit b03b3ae99c
1 changed files with 6 additions and 1 deletions
--- a/colossalai/utils/memory_tracer/async_memtracer.py
+++ b/colossalai/utils/memory_tracer/async_memtracer.py
@ -52,7 +52,12 @@ class AsyncMemoryMonitor:

    def __init__(self, power: int = 10):
        self.keep_measuring = False
-        self.executor = ThreadPoolExecutor(max_workers=1)
+
+        current_device = get_current_device()
+        def _set_cuda_device():
+            torch.cuda.set_device(current_device)
+
+        self.executor = ThreadPoolExecutor(max_workers=1, initializer=_set_cuda_device)
        self.monitor_thread = None
        self.interval = 1 / (10**power)
        self.time_stamps = []