double check optimum

aws · Jun 27, 2023 · 94c484c · 94c484c
1 parent 6a368d4
commit 94c484c
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/sagemaker_huggingface_inference_toolkit/optimum_utils.py b/src/sagemaker_huggingface_inference_toolkit/optimum_utils.py
@@ -17,7 +17,10 @@
 import os
 
 
-_optimum_neuron = importlib.util.find_spec("optimum.neuron") is not None
+_optimum_neuron = False
+if importlib.util.find_spec("optimum") is not None:
+ if importlib.util.find_spec("optimum.neuron") is not None:
+ _optimum_neuron = True
 
 logger = logging.getLogger(__name__)