fixed mem monitor device (#433)

fixed mem monitor device

fixed mem monitor device (#433)
fixed mem monitor device
b03b3ae9 · Frank Lee · GitHub · 14a70942 · b03b3ae9
Unverified Commit b03b3ae9 authored Mar 16, 2022 by Frank Lee Committed by GitHub Mar 16, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 6 additions and 1 deletion

colossalai/utils/memory_tracer/async_memtracer.py colossalai/utils/memory_tracer/async_memtracer.py +6 -1

No files found.
--- 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 = []