[GitHub Runner] Fix flax runner (#13357)

Commit

4 years ago

[GitHub Runner] Fix flax runner (#13357) * correct * also comment out multi-gpu test push