update integration with Tensorflow (#119)

* Apply suggestions from code review

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Jirka Borovec <[email protected]>

Author: 3 people

src/litmodels/io/utils.py

@@ -1,14 +1,31 @@
+import os
 import pickle
+from contextlib import contextmanager
 from pathlib import Path
 from typing import Any, Union
 
 from lightning_utilities import module_available
 from lightning_utilities.core.imports import RequirementCache
 
+
+@contextmanager
+def _suppress_os_stderr():
+    devnull_fd = os.open(os.devnull, os.O_WRONLY)
+    old_stderr_fd = os.dup(2)
+    os.dup2(devnull_fd, 2)  # redirect stderr (fd 2) to /dev/null
+    os.close(devnull_fd)
+    try:
+        yield
+    finally:
+        os.dup2(old_stderr_fd, 2)  # restore stderr
+        os.close(old_stderr_fd)
+
+
 _JOBLIB_AVAILABLE = module_available("joblib")
 _PYTORCH_AVAILABLE = module_available("torch")
-_TENSORFLOW_AVAILABLE = module_available("tensorflow")
-_KERAS_AVAILABLE = RequirementCache("tensorflow >=2.0.0")
+with _suppress_os_stderr():
+    _TENSORFLOW_AVAILABLE = module_available("tensorflow")
+    _KERAS_AVAILABLE = RequirementCache("tensorflow >=2.0.0")
 
 if _JOBLIB_AVAILABLE:
     import joblib