formatting fix

Author: rdyro

torch2jax/api.py

@@ -150,7 +150,7 @@ def torch2jax(
         example_kw: Example keyword arguments. Defaults to None.
         output_shapes: Output shapes or shapes + dtype struct. Defaults to None.
         output_sharding_spec: jax.sharding.PartitionSpec specifying the sharding spec of the output, uses input mesh.
-        vmap_method: batching method, see 
+        vmap_method: batching method, see
             [https://docs.jax.dev/en/latest/ffi.html#batching-with-vmap](https://docs.jax.dev/en/latest/ffi.html#batching-with-vmap)
 
             NOTE: only vmap_method="sequntial" is supported non-experimentally

torch2jax/gradients.py

@@ -46,7 +46,7 @@ def torch2jax_with_vjp(
             library PyTorch code may need this fallback. Defaults to True (i.e., do not use fallback).
         output_sharding_spec: (not supported) sharding spec of the output, use shard_map instead for a device-local
             version of this function
-        vmap_method: batching method, see 
+        vmap_method: batching method, see
             [https://docs.jax.dev/en/latest/ffi.html#batching-with-vmap](https://docs.jax.dev/en/latest/ffi.html#batching-with-vmap)
 
             NOTE: only vmap_method="sequntial" is supported non-experimentally