correctly init the keras backend in the unit tests

aertslab · Dec 2, 2024 · 61c080a · 61c080a
1 parent edab6dd
commit 61c080a
Showing 1 changed file with 18 additions and 0 deletions.
diff --git a/tests/__init__.py b/tests/__init__.py
@@ -1 +1,19 @@
 """Init file for the tests module."""
+
+import os
+
+# Set keras backend
+try:
+    import tensorflow as tf
+
+    os.environ["KERAS_BACKEND"] = "tensorflow"
+    os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+except ImportError:
+    try:
+        import torch
+
+        os.environ["KERAS_BACKEND"] = "torch"
+    except ImportError as e:
+        raise ImportError(
+            "No backend found. Please install either tensorflow or pytorch."
+        ) from e