fix: numba.*_num_threads resets torch num_threads NVIDIA#8141

temporary fix until numba/numba#9387 gets resolved. Signed-off-by: Iztok Lebar Bajec <[email protected]>
itzsimpl · Jan 10, 2024 · 14dae6f · 14dae6f
1 parent 76a712a
commit 14dae6f
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/nemo/collections/asr/parts/numba/rnnt_loss/utils/cpu_utils/cpu_rnnt.py b/nemo/collections/asr/parts/numba/rnnt_loss/utils/cpu_utils/cpu_rnnt.py
@@ -203,10 +203,13 @@ def __init__(
         self.num_threads_ = num_threads
         self.batch_first = batch_first
 
+        _torch_num_threads = torch.get_num_threads()
         if num_threads > 0:
             numba.set_num_threads(min(multiprocessing.cpu_count(), num_threads))
+            self.num_threads_ = numba.get_num_threads()
         else:
             self.num_threads_ = numba.get_num_threads()
+        torch.set_num_threads(_torch_num_threads)
 
     def cost_and_grad_kernel(
         self,

diff --git a/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/gpu_rnnt.py b/nemo/collections/asr/parts/numba/rnnt_loss/utils/cuda_utils/gpu_rnnt.py
@@ -82,11 +82,13 @@ def __init__(
         self.num_threads_ = num_threads
         self.stream_ = stream  # type: cuda.cudadrv.driver.Stream
 
+        _torch_num_threads = torch.get_num_threads()
         if num_threads > 0:
             numba.set_num_threads(min(multiprocessing.cpu_count(), num_threads))
             self.num_threads_ = numba.get_num_threads()
         else:
             self.num_threads_ = numba.get_num_threads()
+        torch.set_num_threads(_torch_num_threads)
 
     def log_softmax(self, acts: torch.Tensor, denom: torch.Tensor):
         """