github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[extension] fixed exception catch (#5342)

pull/5345/head
Frank Lee 2024-01-31 18:09:49 +08:00 committed by GitHub
parent 71321a07cf
commit abd8e77ad8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
extensions/cpp_extension.py
View File

@ -126,7 +126,7 @@ class _CppExtension(_Extension):
def load(self): def load(self):
try: try:
op_kernel = self.import_op() op_kernel = self.import_op()
except ImportError: except (ImportError, ModuleNotFoundError):
# if import error occurs, it means that the kernel is not pre-built # if import error occurs, it means that the kernel is not pre-built
# so we build it jit # so we build it jit
op_kernel = self.build_jit() op_kernel = self.build_jit()