mirror of https://github.com/hpcaitech/ColossalAI
[extension] fixed exception catch (#5342)
parent
71321a07cf
commit
abd8e77ad8
|
@ -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()
|
||||
|
|
Loading…
Reference in New Issue