Use optional import

Project-MONAI · Nov 29, 2023 · 4caafe0 · 4caafe0
1 parent a6bb925
commit 4caafe0
Showing 1 changed file with 3 additions and 14 deletions.
diff --git a/monai/networks/nets/autoencoderkl.py b/monai/networks/nets/autoencoderkl.py
@@ -11,7 +11,6 @@
 
 from __future__ import annotations
 
-import importlib.util
 import math
 from collections.abc import Sequence
 from typing import List
@@ -21,21 +20,11 @@
 import torch.nn.functional as F
 
 from monai.networks.blocks import Convolution
-from monai.utils import ensure_tuple_rep
 
 # To install xformers, use pip install xformers==0.0.16rc401
-if importlib.util.find_spec("xformers") is not None:
- import xformers
- import xformers.ops
-
- has_xformers = True
-else:
- xformers = None
- has_xformers = False
-
-# TODO: Use MONAI's optional_import
-# from monai.utils import optional_import
-# xformers, has_xformers = optional_import("xformers.ops", name="xformers")
+from monai.utils import ensure_tuple_rep, optional_import
+
+xformers, has_xformers = optional_import("xformers")
 
 __all__ = ["AutoencoderKL"]