[math] Remove the logs that taichi.init() print (#609)

brainpy · Feb 1, 2024 · a9996f3 · a9996f3
1 parent bde7f8a
commit a9996f3
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/brainpy/_src/math/op_register/taichi_aot_based.py b/brainpy/_src/math/op_register/taichi_aot_based.py
@@ -1,5 +1,7 @@
+import contextlib
 import hashlib
 import inspect
+import io
 import os
 import pathlib
 import platform
@@ -173,8 +175,8 @@ def _build_kernel(
     arch = ti.cuda
   else:
     raise ValueError(f'Unknown device: {device}')
-
-  ti.init(arch=arch)
+  with contextlib.redirect_stdout(io.StringIO()):
+    ti.init(arch=arch)
 
   # check arch is available
   if ti.lang.impl.current_cfg().arch != arch: