run yapf

sail-sg · Aug 12, 2024 · 18e6ac2 · 18e6ac2
1 parent e69ebb8
commit 18e6ac2
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/envpool/python/xla_template.py b/envpool/python/xla_template.py
@@ -91,12 +91,8 @@ def translation(c: Any, *args: Any, platform: str = "cpu") -> Any:
   prim.multiple_results = (len(out_specs) > 1)
   prim.def_impl(partial(xla.apply_primitive, prim))
   prim.def_abstract_eval(abstract)
-  interpreters.mlir["cpu"][prim] = partial(
-    translation, platform="cpu"
-  )
-  interpreters.mlir["gpu"][prim] = partial(
-    translation, platform="gpu"
-  )
+  interpreters.mlir["cpu"][prim] = partial(translation, platform="cpu")
+  interpreters.mlir["gpu"][prim] = partial(translation, platform="gpu")
 
   def call(*args: Any) -> Any:
     return prim.bind(*args)