#mjx don't specify JAX default backend

PiperOrigin-RevId: 785472058 Change-Id: Iba2d3e87dd2d78e6eee8fb579a760250a91b5d23
2025-07-21 09:46:49 -07:00
parent fc13995dd4
commit 921bca5105
1 changed files with 1 additions and 1 deletions
@@ -148,7 +148,7 @@ def _resolve_impl_and_device(
  elif (has_impl, has_device) == (False, True):
    impl = _resolve_impl(device)
  else:
-    device = jax.devices(jax.default_backend())[0]
+    device = jax.devices()[0]
    logging.info('Using JAX default device: %s.', device)
    impl = _resolve_impl(device)