Skip to content

Commit

Permalink
fix: removing .encode() in call to register_custom_call_target inside…
Browse files Browse the repository at this point in the history
… _make_xla_function
  • Loading branch information
Cyprien Courtot committed Jul 16, 2024
1 parent f411fc2 commit 9b41bce
Showing 1 changed file with 2 additions and 4 deletions.
6 changes: 2 additions & 4 deletions envpool/python/xla_template.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,12 +53,10 @@ def _make_xla_function(
out_specs = _normalize_specs(out_specs)
cpu_capsule, gpu_capsule = capsules
xla_client.register_custom_call_target(
f"{type(obj).__name__}_{id(obj)}_{name}_cpu".encode(),
cpu_capsule,
platform="cpu"
f"{type(obj).__name__}_{id(obj)}_{name}_cpu", cpu_capsule, platform="cpu"
)
xla_client.register_custom_call_target(
f"{type(obj).__name__}_{id(obj)}_{name}_gpu".encode(),
f"{type(obj).__name__}_{id(obj)}_{name}_gpu",
gpu_capsule,
platform="gpu",
)
Expand Down

0 comments on commit 9b41bce

Please sign in to comment.