From ae139513c05553e0646c8011e5903cf2decd15c0 Mon Sep 17 00:00:00 2001 From: Sam Duffield Date: Wed, 28 Feb 2024 17:47:48 +0000 Subject: [PATCH] Remove import warning from torchopt --- pyproject.toml | 2 +- uqlib/__init__.py | 15 +++++++++++++-- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b8c0e5e0..06d38b2a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,4 +24,4 @@ packages = ["uqlib"] [tool.ruff] [tool.ruff.lint.per-file-ignores] -"__init__.py" = ["F401", "F821"] \ No newline at end of file +"__init__.py" = ["F401", "F821", "E402"] \ No newline at end of file diff --git a/uqlib/__init__.py b/uqlib/__init__.py index 3d7f35cd..a142c82c 100644 --- 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