feat(library): allow fallback when kernel fails

huggingface · Feb 8, 2024 · 6f297be · 6f297be
1 parent db5928b
commit 6f297be
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/quanto/library/ops.py b/quanto/library/ops.py
@@ -1,3 +1,4 @@
+import warnings
 from contextlib import contextmanager
 
 import torch
@@ -41,7 +42,13 @@ def define(name, schema):
     @torch.library.impl(f"quanto::{name}", "default")
     def impl(*args, **kwargs):
         if _ext_enabled:
-            return getattr(torch.ops.quanto_ext, name)(*args, **kwargs)
+            try:
+                return getattr(torch.ops.quanto_ext, name)(*args, **kwargs)
+            except Exception as e:
+                warnings.warn(
+                    f"A {type(e)} exception occured while calling the optimized kernel for quanto::{name}."
+                    "Falling back to default implementation."
+                )
         return getattr(torch.ops.quanto_py, name)(*args, **kwargs)