add conditional import for tensorflow

NOEL-MNI · Sep 14, 2023 · 217907a · 217907a
1 parent b6006e7
commit 217907a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/app/utils/image_processing.py b/app/utils/image_processing.py
@@ -8,8 +8,8 @@
 import multiprocessing
 import numpy as np
 # import zipfile
-
-from antspynet.utilities import brain_extraction
+if os.environ.get("BRAIN_MASKING") == "cpu":
+    from antspynet.utilities import brain_extraction
 from .deepmask import *
 from .helpers import *
 from matplotlib.backends.backend_pdf import PdfPages