[extension] fixed exception catch (#5342)

10 months ago · abd8e77ad8
parent 71321a07cf
commit abd8e77ad8
1 changed files with 1 additions and 1 deletions
--- a/extensions/cpp_extension.py
+++ b/extensions/cpp_extension.py
@ -126,7 +126,7 @@ class _CppExtension(_Extension):
    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()