Pass torch.Tensor to jax.dlpack.from_dlpack directly

The code path accepting a DLPack capsule is deprecated and will be removed soon.

PiperOrigin-RevId: 797799893
Change-Id: I54a2bdad6eb9ee8efaa230a8cb416b7179f67fe2

Author: superbobry

brax/io/torch.py

@@ -52,9 +52,7 @@ def torch_to_jax(value: Any) -> Any:
 @torch_to_jax.register(torch.Tensor)
 def _tensor_to_jax(value: torch.Tensor) -> jax.Array:
   """Converts a PyTorch Tensor into a jax.Array."""
-  tensor = torch_dlpack.to_dlpack(value)
-  tensor = jax_dlpack.from_dlpack(tensor)
-  return tensor
+  return jax_dlpack.from_dlpack(value)
 
 
 @torch_to_jax.register(abc.Mapping)