[flax:examples:sst2] Fix pytype errors.#5118

copybara-service[bot] merged 1 commit intomaingoogle/flax:mainfrom

test_838871691google/flax:test_838871691Copy head branch name to clipboard