use register_custom_call_target in stead of register_cpu_custom_call_…

…target
flaport · Sep 19, 2023 · 1d5cf72 · 1d5cf72
1 parent ee67207
commit 1d5cf72
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/klujax.py b/klujax.py
@@ -57,7 +57,7 @@ def xla_register_cpu(primitive, cpp_fun):
     name = primitive.name.encode()
 
     def decorator(fun):
-        xla_client.register_cpu_custom_call_target(
+        xla_client.register_custom_call_target(
             name,
             cpp_fun(),
         )