Use custom docker image to run code and remove jax-tpu-emb wheel#152

ajkv-google merged 1 commit intomainAI-Hypercomputer/RecML:mainfrom

ajkv/docker-updateAI-Hypercomputer/RecML:ajkv/docker-updateCopy head branch name to clipboard