Merge pull request #27 from normal-computing/clean-import

Remove import warning from torchopt
normal-computing · Feb 28, 2024 · e396cd5 · e396cd5
2 parents cb47f7d + ae13951
commit e396cd5
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -24,4 +24,4 @@ packages = ["uqlib"]
 
 [tool.ruff]
 [tool.ruff.lint.per-file-ignores]
-"__init__.py" = ["F401", "F821"]
+"__init__.py" = ["F401", "F821", "E402"]
diff --git a/uqlib/__init__.py b/uqlib/__init__.py
@@ -2,9 +2,8 @@
 from uqlib import laplace
 from uqlib import sgmcmc
 from uqlib import types
-from uqlib import vi
 from uqlib import optim
-from uqlib import torchopt
+
 
 from uqlib.utils import model_to_function
 from uqlib.utils import linearized_forward_diag
@@ -23,3 +22,15 @@
 from uqlib.utils import tree_map_inplacify_
 from uqlib.utils import flexi_tree_map
 from uqlib.utils import per_samplify
+
+
+import logging
+
+logger = logging.getLogger("torch.distributed.elastic.multiprocessing.redirects")
+logger.setLevel(logging.ERROR)
+
+from uqlib import vi
+from uqlib import torchopt
+
+del logging
+del logger