some update (#2738)

intel-analytics · Sep 26, 2021 · d99af71 · d99af71
1 parent de4a1cf
commit d99af71
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/python/orca/src/bigdl/orca/torch/torch_loss.py b/python/orca/src/bigdl/orca/torch/torch_loss.py
@@ -16,11 +16,15 @@
 import sys
 from bigdl.nn.criterion import Criterion
 from pyspark.serializers import CloudPickleSerializer
+from importlib.util import find_spec
 
 if sys.version < '3.7':
     print("WARN: detect python < 3.7, if you meet zlib not available " +
           "exception on yarn, please update your python to 3.7")
 
+if find_spec('jep') is None:
+    raise Exception("jep not found, please install jep first.")
+
 
 class TorchLoss(Criterion):
     """

diff --git a/python/orca/src/bigdl/orca/torch/torch_model.py b/python/orca/src/bigdl/orca/torch/torch_model.py
@@ -22,11 +22,15 @@
 from zoo.common.utils import callZooFunc
 from pyspark.serializers import CloudPickleSerializer
 from zoo.pipeline.api.torch.utils import trainable_param
+from importlib.util import find_spec
 
 if sys.version < '3.7':
     print("WARN: detect python < 3.7, if you meet zlib not available " +
           "exception on yarn, please update your python to 3.7")
 
+if find_spec('jep') is None:
+    raise Exception("jep not found, please install jep first.")
+
 
 class TorchModel(Layer):
     """