feat(library): allow fallback when kernel fails

huggingface · Feb 8, 2024 · cc29588 · cc29588
1 parent 4bffe4e
commit cc29588
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)