Skip to content

Commit

Permalink
Update on "[Executorch][BE] Rename sdpa_with_kv_cache.py to custom_op…
Browse files Browse the repository at this point in the history
…s.py"

Because now we have more than sdpa_with_kv_cache in it

Differential Revision: [D66269486](https://our.internmc.facebook.com/intern/diff/D66269486/)

[ghstack-poisoned]
  • Loading branch information
kimishpatel committed Nov 22, 2024
2 parents 863b55c + 9515add commit a2e5433
Showing 1 changed file with 3 additions and 1 deletion.
4 changes: 3 additions & 1 deletion extension/llm/custom_ops/custom_ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,9 @@
op2 = torch.ops.llama.fast_hadamard_transform.default
assert op2 is not None
except:
libs = list(Path(__file__).parent.resolve().glob("libcustom_ops_aot_lib.*"))
path = Path(__file__).parent.resolve()
logging.info(f"Looking for libcustom_ops_aot_lib.so in {path}")
libs = list(path.glob("libcustom_ops_aot_lib.*"))
assert len(libs) == 1, f"Expected 1 library but got {len(libs)}"
logging.info(f"Loading custom ops library: {libs[0]}")
torch.ops.load_library(libs[0])
Expand Down

0 comments on commit a2e5433

Please sign in to comment.