diff --git a/extensions/cpp_extension.py b/extensions/cpp_extension.py index b4c40c9f1105..3adb65fb8f4e 100644 --- a/extensions/cpp_extension.py +++ b/extensions/cpp_extension.py @@ -126,7 +126,7 @@ def cxx_flags(self) -> List[str]: def load(self): try: op_kernel = self.import_op() - except ImportError: + except (ImportError, ModuleNotFoundError): # if import error occurs, it means that the kernel is not pre-built # so we build it jit op_kernel = self.build_jit()