fix(jax): avoid device lookup during type embedding tracing#5736

wanghan-iapcm merged 1 commit into

deepmodeling:masterdeepmodeling/deepmd-kit:masterfrom

njzjz:fix/jax-type-embed-tracingnjzjz/deepmd-kit:fix/jax-type-embed-tracingCopy head branch name to clipboard