Fix PPO input sharding to match pmap axis - #674

Fix PPO input sharding to match pmap axis#674

lmj0209 wants to merge 1 commit into

google:maingoogle/brax:mainfrom

lmj0209:fix/ppo-jax-device-putlmj0209/brax:fix/ppo-jax-device-putCopy head branch name to clipboard